From bd397871030d9f12d360fafd32dc15957e655543 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 21 Jul 2026 19:11:11 +0800 Subject: [PATCH 01/10] feat(dpa4c): add compact invariant descriptor and CUDA deployment Introduce DPA4C as a graph-native descriptor built from degree-scaled Cartesian moments, exact invariant readouts, and pair-conditioned radial modes. - support backend-neutral training, serialization, graph export, calibration, and mixed-precision execution - add compressed CUDA and canonical inference for the supported channel and angular profiles - expose neighborhood masses, remove the fixed-capacity path, and tile compressed evaluation to bound memory - cover parity, gradients, compression, export, and end-to-end energy/force/virial behavior --- deepmd/dpmodel/descriptor/__init__.py | 4 + deepmd/dpmodel/descriptor/dpa4_nn/__init__.py | 6 + .../dpmodel/descriptor/dpa4_nn/embedding.py | 25 +- deepmd/dpmodel/descriptor/dpa4_nn/mlp.py | 204 +++ deepmd/dpmodel/descriptor/dpa4_nn/radial.py | 21 +- deepmd/dpmodel/descriptor/dpa4c.py | 1429 ++++++++++++++++ .../dpmodel/descriptor/dpa4c_nn/__init__.py | 38 + .../dpmodel/descriptor/dpa4c_nn/bispectrum.py | 315 ++++ .../dpmodel/descriptor/dpa4c_nn/geometry.py | 271 +++ .../dpmodel/descriptor/dpa4c_nn/pair_film.py | 196 +++ deepmd/dpmodel/descriptor/dpa4c_nn/readout.py | 527 ++++++ deepmd/dpmodel/loss/ener.py | 18 +- deepmd/kernels/cuda/__init__.py | 7 +- deepmd/kernels/cuda/dpa1/canonical.py | 58 +- deepmd/kernels/cuda/dpa1/graph_compress.py | 60 +- .../kernels/cuda/dpa1/graph_energy_force.py | 32 +- deepmd/kernels/cuda/dpa4c/__init__.py | 28 + deepmd/kernels/cuda/dpa4c/canonical.py | 418 +++++ deepmd/kernels/cuda/dpa4c/graph_compress.py | 1395 ++++++++++++++++ deepmd/kernels/cuda/edge_force_virial.py | 67 +- deepmd/kernels/cuda/graph_fitting.py | 348 +++- deepmd/pt_expt/descriptor/__init__.py | 4 + deepmd/pt_expt/descriptor/dpa4c.py | 529 ++++++ deepmd/pt_expt/entrypoints/main.py | 17 +- deepmd/pt_expt/infer/deep_eval.py | 37 +- deepmd/pt_expt/model/ener_model.py | 105 +- deepmd/pt_expt/model/graph_lower.py | 6 +- deepmd/pt_expt/utils/canonical_graph.py | 37 +- deepmd/pt_expt/utils/network.py | 8 +- deepmd/pt_expt/utils/serialization.py | 56 +- deepmd/utils/argcheck.py | 115 ++ doc/model/dpa4c.md | 334 ++++ doc/model/index.rst | 1 + examples/water/dpa4c/README.md | 53 + examples/water/dpa4c/input.json | 78 + source/api_c/include/c_api.h | 6 +- source/api_c/include/deepmd.hpp | 4 +- source/api_c/src/c_api.cc | 4 +- source/api_cc/include/DeepPot.h | 9 +- source/api_cc/include/DeepPotPTExpt.h | 8 +- source/api_cc/include/commonPT.h | 18 +- source/api_cc/src/DeepPot.cc | 8 +- source/api_cc/src/DeepPotPTExpt.cc | 30 +- .../api_cc/tests/test_neighbor_list_data.cc | 6 +- source/lmp/pair_deepmd_kokkos.cpp | 188 ++- source/lmp/pair_deepmd_kokkos.h | 19 +- source/op/pt/CMakeLists.txt | 19 +- source/op/pt/dpa1_graph_compress.cu | 42 +- source/op/pt/dpa1_graph_compress_kernel.cuh | 7 + source/op/pt/dpa1_graph_compress_launch.h | 1 + source/op/pt/dpa1_graph_energy_force.cu | 12 +- source/op/pt/dpa4c_graph_compress.cu | 794 +++++++++ source/op/pt/dpa4c_graph_compress.cuh | 635 +++++++ source/op/pt/dpa4c_graph_compress_c128.cu | 12 + source/op/pt/dpa4c_graph_compress_c16.cu | 12 + source/op/pt/dpa4c_graph_compress_c32.cu | 12 + source/op/pt/dpa4c_graph_compress_c64.cu | 12 + source/op/pt/dpa4c_graph_compress_c8.cu | 12 + source/op/pt/dpa4c_graph_compress_kernel.cuh | 1463 +++++++++++++++++ source/op/pt/dpa4c_graph_compress_launch.h | 274 +++ source/op/pt/edge_force_virial.cu | 176 +- source/op/pt/graph_fitting.cu | 552 +++++-- source/op/pt/graph_ops.h | 57 +- .../common/dpmodel/test_descriptor_dpa4c.py | 756 +++++++++ source/tests/common/test_examples.py | 1 + .../pt_expt/descriptor/test_dpa1_cuda.py | 37 +- source/tests/pt_expt/descriptor/test_dpa4c.py | 428 +++++ .../pt_expt/descriptor/test_dpa4c_cuda.py | 1059 ++++++++++++ .../pt_expt/model/test_dpa4c_graph_lower.py | 212 +++ .../pt_expt/utils/test_canonical_graph.py | 30 +- 70 files changed, 13156 insertions(+), 606 deletions(-) create mode 100644 deepmd/dpmodel/descriptor/dpa4_nn/mlp.py create mode 100644 deepmd/dpmodel/descriptor/dpa4c.py create mode 100644 deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py create mode 100644 deepmd/dpmodel/descriptor/dpa4c_nn/bispectrum.py create mode 100644 deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py create mode 100644 deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py create mode 100644 deepmd/dpmodel/descriptor/dpa4c_nn/readout.py create mode 100644 deepmd/kernels/cuda/dpa4c/__init__.py create mode 100644 deepmd/kernels/cuda/dpa4c/canonical.py create mode 100644 deepmd/kernels/cuda/dpa4c/graph_compress.py create mode 100644 deepmd/pt_expt/descriptor/dpa4c.py create mode 100644 doc/model/dpa4c.md create mode 100644 examples/water/dpa4c/README.md create mode 100644 examples/water/dpa4c/input.json create mode 100644 source/op/pt/dpa4c_graph_compress.cu create mode 100644 source/op/pt/dpa4c_graph_compress.cuh create mode 100644 source/op/pt/dpa4c_graph_compress_c128.cu create mode 100644 source/op/pt/dpa4c_graph_compress_c16.cu create mode 100644 source/op/pt/dpa4c_graph_compress_c32.cu create mode 100644 source/op/pt/dpa4c_graph_compress_c64.cu create mode 100644 source/op/pt/dpa4c_graph_compress_c8.cu create mode 100644 source/op/pt/dpa4c_graph_compress_kernel.cuh create mode 100644 source/op/pt/dpa4c_graph_compress_launch.h create mode 100644 source/tests/common/dpmodel/test_descriptor_dpa4c.py create mode 100644 source/tests/pt_expt/descriptor/test_dpa4c.py create mode 100644 source/tests/pt_expt/descriptor/test_dpa4c_cuda.py create mode 100644 source/tests/pt_expt/model/test_dpa4c_graph_lower.py diff --git a/deepmd/dpmodel/descriptor/__init__.py b/deepmd/dpmodel/descriptor/__init__.py index 0b3570b4fe..ae9cc66e39 100644 --- a/deepmd/dpmodel/descriptor/__init__.py +++ b/deepmd/dpmodel/descriptor/__init__.py @@ -11,6 +11,9 @@ from .dpa4 import ( DescrptDPA4, ) +from .dpa4c import ( + DescrptDPA4C, +) from .hybrid import ( DescrptHybrid, ) @@ -38,6 +41,7 @@ "DescrptDPA2", "DescrptDPA3", "DescrptDPA4", + "DescrptDPA4C", "DescrptHybrid", "DescrptSeA", "DescrptSeAttenV2", diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py index 85c854e69c..817b50c841 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/__init__.py @@ -73,6 +73,10 @@ merge_lora_into_base, strip_lora_from_extra_state, ) +from .mlp import ( + SwiGLUMLP, + resolve_swiglu_hidden_width, +) from .norm import ( EquivariantRMSNorm, ReducedEquivariantRMSNorm, @@ -159,6 +163,7 @@ "SeZMTypeEmbedding", "SpinEmbedding", "SwiGLU", + "SwiGLUMLP", "WignerDCalculator", "apply_lora_to_sezm", "build_cartesian_basis", @@ -189,6 +194,7 @@ "quaternion_z_rotation", "resolve_s2_grid_resolution", "resolve_so3_grid", + "resolve_swiglu_hidden_width", "safe_norm", "segment_envelope_gated_softmax", "so3_packed_index", diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py index 82c55dd8c4..be6bb291b8 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py @@ -128,26 +128,41 @@ def __init__( # === Step 2. Register the embedding table parameter === self.adam_type_embedding = table - def call(self, atype: Any) -> Any: + def call(self, atype: Any | None = None) -> Any: """ Gather type embeddings. Parameters ---------- atype - Atom types with shape (...,). Valid type range is [0, ntypes-1]. + Atom types with shape (...). Valid type range is [0, ntypes-1]. + If omitted, return the complete embedding table, including the + optional padding row. This form is used by graph-native descriptor + ABIs that precompute the table once per forward call. Returns ------- Array - Type embeddings with shape (..., embed_dim). + Gathered type embeddings with shape ``(..., embed_dim)`` when + ``atype`` is provided. Otherwise, the complete table with shape + ``(ntypes + int(padding), embed_dim)``. """ + # === Step 1. Return the complete graph-native lookup table === + if atype is None: + xp = array_api_compat.array_namespace(self.adam_type_embedding) + return xp_asarray_nodetach( + xp, + self.adam_type_embedding[...], + device=array_api_compat.device(self.adam_type_embedding), + ) + + # === Step 2. Gather rows for an explicit atom-type tensor === xp = array_api_compat.array_namespace(atype) weight = xp_asarray_nodetach( xp, self.adam_type_embedding[...], device=array_api_compat.device(atype) ) - # torch.embedding gather: flatten the indices to int64, take the rows, - # then restore the original index shape. + # Flattening provides one backend-neutral gather while preserving every + # leading batch or graph dimension on restoration. index = xp.astype(xp.reshape(atype, (-1,)), xp.int64) if self.padding: index = remap_atype_to_padding(index, self.ntypes + 1) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/mlp.py b/deepmd/dpmodel/descriptor/dpa4_nn/mlp.py new file mode 100644 index 0000000000..da90d08117 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/mlp.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bias-free SwiGLU multilayer perceptrons for DPA4-family descriptors.""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.utils.network import ( + NativeLayer, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .activation import ( + SwiGLU, +) + + +def resolve_swiglu_hidden_width(width: int, multiple: int = 8) -> int: + r"""Return the parameter-matched SwiGLU hidden width. + + The post-gate width is :math:`8d/3`, rounded up to ``multiple``. The + corresponding hidden affine map produces twice this width for the value + and gate branches. + + Parameters + ---------- + width + Input and output model width. + multiple + Alignment multiple for the post-gate hidden width. + + Returns + ------- + int + Aligned post-gate hidden width. + """ + if width <= 0: + raise ValueError(f"`width` must be positive, got {width}") + if multiple <= 0: + raise ValueError(f"`multiple` must be positive, got {multiple}") + numerator = 8 * int(width) + denominator = 3 * int(multiple) + return int(multiple) * ((numerator + denominator - 1) // denominator) + + +class SwiGLUMLP(NativeOP): + """Apply bias-free SwiGLU hidden layers and a linear output projection. + + For hidden width ``H``, each hidden affine map produces ``2H`` channels. + :class:`SwiGLU` splits them into equal gate and value branches and returns + ``SiLU(gate) * value`` with width ``H``. The final layer is linear. + + Parameters + ---------- + mlp_layers + Layer widths including input, hidden, and output dimensions. + output_scale + Fixed multiplier applied to the final output. + precision + Parameter precision. + trainable + Whether the linear weights are trainable. + seed + Random seed. + """ + + def __init__( + self, + mlp_layers: list[int], + *, + output_scale: float = 1.0, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + if len(mlp_layers) < 2: + raise ValueError("`mlp_layers` must contain input and output widths") + if any(width <= 0 for width in mlp_layers): + raise ValueError(f"`mlp_layers` must be positive, got {mlp_layers}") + self.mlp_layers = [int(width) for width in mlp_layers] + self.output_scale = float(output_scale) + self.precision = str(precision) + self.trainable = bool(trainable) + + layers = [] + for index, (width_in, width_out) in enumerate( + zip(self.mlp_layers[:-1], self.mlp_layers[1:], strict=True) + ): + is_output = index == len(self.mlp_layers) - 2 + layers.append( + NativeLayer( + width_in, + width_out if is_output else 2 * width_out, + bias=False, + precision=self.precision, + seed=child_seed(seed, index), + trainable=self.trainable, + ) + ) + self.layers = layers + self.activation = SwiGLU() + + def call(self, inputs: Any) -> Any: + """Evaluate the SwiGLU MLP. + + Parameters + ---------- + inputs + Input with shape ``(..., mlp_layers[0])``. + + Returns + ------- + Any + Output with shape ``(..., mlp_layers[-1])``. + """ + return self.call_output(self.call_hidden(inputs)) + + def call_hidden(self, inputs: Any) -> Any: + """Evaluate every hidden layer and return the latent state. + + The latent state is exposed separately so that several output heads + can branch off one trunk evaluation. + + Parameters + ---------- + inputs + Input with shape ``(..., mlp_layers[0])``. + + Returns + ------- + Any + Activated latent state with shape ``(..., mlp_layers[-2])``. + """ + output = inputs + for layer in self.layers[:-1]: + output = self.activation(layer(output)) + return output + + def call_output(self, hidden: Any) -> Any: + """Apply the final scaled linear projection to a latent state. + + Parameters + ---------- + hidden + Latent state with shape ``(..., mlp_layers[-2])``, as returned by + :meth:`call_hidden`. + + Returns + ------- + Any + Output with shape ``(..., mlp_layers[-1])``. + """ + return self.layers[-1](hidden) * self.output_scale + + def serialize(self) -> dict[str, Any]: + """Serialize the MLP configuration and linear weights.""" + return { + "@class": "SwiGLUMLP", + "@version": 1, + "mlp_layers": self.mlp_layers.copy(), + "output_scale": self.output_scale, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + "@variables": { + f"{index}.matrix": to_numpy_array(layer.w) + for index, layer in enumerate(self.layers) + }, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> SwiGLUMLP: + """Deserialize a :class:`SwiGLUMLP`.""" + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + if data.pop("@class") != "SwiGLUMLP": + raise ValueError("Invalid serialized class for SwiGLUMLP") + variables = data.pop("@variables") + obj = cls(**data) + dtype = PRECISION_DICT[obj.precision] + for key, value in variables.items(): + index, _, name = key.partition(".") + if name != "matrix": + raise ValueError(f"Invalid SwiGLUMLP variable {key!r}") + obj.layers[int(index)].w = np.asarray(value, dtype=dtype) + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py index 5ef1c92771..aada21eda5 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py @@ -465,6 +465,11 @@ class RadialBasis(NativeOP): Floating-point precision for the radial basis frequencies and outputs. exponent : int, optional Exponent for the C^3 cutoff envelope polynomial. Default is 7. + apply_envelope : bool, optional + Whether :meth:`call` multiplies the raw basis by the C³ envelope. + The default ``True`` preserves the DPA4 radial contract. Consumers that + apply one shared envelope after combining radial and type features may + request the raw basis with ``False``. """ def __init__( @@ -474,6 +479,7 @@ def __init__( n_radial: int = 10, precision: str = DEFAULT_PRECISION, exponent: int = 7, + apply_envelope: bool = True, ) -> None: self.rcut = float(rcut) if self.rcut <= 0.0: @@ -486,6 +492,7 @@ def __init__( raise ValueError("`basis_type` must be either 'bessel' or 'gaussian'") self.precision = precision self.exponent = int(exponent) + self.apply_envelope = bool(apply_envelope) prec = PRECISION_DICT[self.precision.lower()] self.pi_tensor = math.pi @@ -517,8 +524,9 @@ def call(self, r: Any) -> Any: Returns ------- Array - Radial basis multiplied by C^3 cutoff envelope with shape (N, n_rbf). - The output is smoothly truncated to zero at r = rcut. + Radial basis with shape ``(N, n_radial)``. When + ``apply_envelope=True``, the output includes the C³ envelope and + vanishes smoothly at ``rcut``; otherwise it is the raw basis. """ xp = array_api_compat.array_namespace(r) freqs = xp_asarray_nodetach( @@ -542,9 +550,10 @@ def call(self, r: Any) -> Any: dr = r - freqs # (N, n_rbf) raw = xp.exp(dr * dr * self.gaussian_coeff) # (N, n_rbf) - # === Step 2. Apply C^3 envelope for smooth cutoff === - envelope = self.envelope(r) # (N, 1) - return raw * envelope + # === Step 2. Apply the optional C³ envelope === + if self.apply_envelope: + return raw * self.envelope(r) + return raw def serialize(self) -> dict[str, Any]: """Serialize RadialBasis including trainable frequencies.""" @@ -556,6 +565,7 @@ def serialize(self) -> dict[str, Any]: "basis_type": self.basis_type, "n_radial": self.n_radial, "exponent": self.exponent, + "apply_envelope": self.apply_envelope, "precision": np.dtype(PRECISION_DICT[self.precision]).name, }, "@variables": {"adam_freqs": to_numpy_array(self.adam_freqs)}, @@ -578,6 +588,7 @@ def deserialize(cls, data: dict[str, Any]) -> RadialBasis: n_radial=int(config["n_radial"]), basis_type=str(config.get("basis_type", "bessel")), exponent=int(config.get("exponent", 7)), + apply_envelope=bool(config.get("apply_envelope", True)), precision=precision, ) if variables is not None: diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py new file mode 100644 index 0000000000..4f13407ef8 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -0,0 +1,1429 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Compact and Compressible degree-wise DPA4 descriptor. + +DPA4C is the compact and compressible degree-wise member of the DPA4 family, a +graph-native local descriptor for lightweight training and for distillation +from DPA4. It reuses the DPA4 radial basis, bias-free radial network, type +embedding, and C³ cutoff envelope, and replaces equivariant message passing +with center-local moment reductions. + +Degree :math:`\ell` retains :math:`C_\ell` channels from a profile derived +entirely from the scalar width, for :math:`2\le L\le 4`. Exact degree Grams +preserve the quadratic channel information, and low-rank Cartesian bispectrum +probes couple every non-scalar degree triple allowed by the O(3) triangle and +parity rules. + +Evaluation is a single destination-local scan: one payload carries both +envelope masses and every degree-wise moment. The descriptor therefore +constructs no neighbor pairs, Wigner rotations, or source-node features, and +reduces the edge axis exactly once. +""" + +from __future__ import ( + annotations, +) + +import dataclasses +import math +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.array_api import ( + xp_asarray_nodetach, +) +from deepmd.dpmodel.common import ( + cast_precision, + get_xp_precision, + to_numpy_array, +) +from deepmd.dpmodel.utils import ( + PairExcludeMask, +) +from deepmd.dpmodel.utils.network import ( + NativeLayer, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.dpmodel.utils.update_sel import ( + UpdateSel, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .base_descriptor import ( + BaseDescriptor, +) +from .dpa4_nn import ( + C3CutoffEnvelope, + RadialBasis, + SeZMTypeEmbedding, + SwiGLUMLP, + resolve_swiglu_hidden_width, +) +from .dpa4c_nn import ( + InvariantReadout, + OrderedPairFiLM, + build_angular_basis, + build_moment_indices, + derive_bispectrum_ranks, + derive_degree_channels, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.utils.data_system import ( + DeepmdDataSystem, + ) + from deepmd.utils.path import ( + DPPath, + ) + + +@BaseDescriptor.register("dpa4c") +class DescrptDPA4C(NativeOP, BaseDescriptor): + r"""Construct the Compact and Compressible degree-wise DPA4 descriptor. + + Let :math:`\chi(r)` denote the fixed DPA4 C³ envelope and + :math:`\phi_{ijc}` the ordered type-conditioned radial amplitude, which + already carries one envelope factor. Degree zero is an additive reduction + under one envelope, while every non-scalar degree carries a second + envelope factor and its own matched normalizer: + + .. math:: + + d^{(0)}_i=\sum_j\chi_{ij}^2,\qquad + d^{(+)}_i=\sum_j\chi_{ij}^4,\qquad + n^{(\bullet)}_i=\bigl(d^{(\bullet)}_i+\tfrac14\bigr)^{-1/2},\\ + X^{(0)}_{ic}=n^{(0)}_i\sum_j\phi_{ijc},\qquad + X^{(\ell)}_{ic} + =n^{(+)}_i\sum_j\chi_{ij}\phi^{(\ell)}_{ijc} + B^{(\ell)}(\hat{\mathbf r}_{ij}). + + Degree :math:`\ell` reads the leading :math:`C_\ell` channels of the one + shared radial map, so the tabulated edge width equals :math:`C_0`. The + invariant readout then contracts these moments into exact channel Grams, a + fixed O(3)-even Cartesian bispectrum over low-rank probes, and the + projected quartic. + + Two further blocks close the output. The divisors + :math:`1/n^{(0)}` and :math:`1/n^{(+)}` are emitted alongside the + invariants, because normalization is otherwise irreversible and the + effective coordination they encode would reach neither the readout nor the + fitting network. The center type embedding is concatenated last as an + independent block. + + Parameters + ---------- + rcut + Outer cutoff radius in Å. + ntypes + Number of atom types. + channels + Scalar degree-zero and edge-amplitude width. Supported values are 8, + 16, 32, 64, and 128. + lmax + Maximum angular degree. Supported values are 2, 3, and 4. + basis_type + DPA4 radial basis type: ``"bessel"`` or ``"gaussian"``. + n_radial + Number of DPA4 radial basis functions forming the fixed analytic + radial input. + radial_modes + Number :math:`R` of shared radial mode profiles that every ordered + atom-type pair mixes with its own coefficients. Zero leaves each pair + with a rescaled copy of one shared radial function. + use_amp + Whether the per-edge stage runs under bfloat16 automatic mixed + precision on CUDA during training. This is an execution policy rather + than model state: it is never serialized, the backend-neutral + equations only record it, and the autocast region itself is a backend + concern. Evaluation and inference follow ``DP_AMP_INFER`` instead. + exclude_types + Ordered atom-type pairs excluded from the descriptor. + precision + Floating-point precision of descriptor parameters. + trainable + Whether descriptor parameters are trainable. + type_map + Atom-type names. + seed + Random seed. + spin + Reserved for descriptor API compatibility; only ``None`` is supported. + + Raises + ------ + TypeError + If ``channels`` or ``lmax`` is not an integer. + ValueError + If ``rcut``, ``ntypes``, or ``n_radial`` is not positive, if + ``radial_modes`` is negative, or if ``channels`` or ``lmax`` is + outside its supported set. + NotImplementedError + If ``spin`` is given. + """ + + _update_sel_cls = UpdateSel + _ENVELOPE_EXPONENT = 5 + _DEGREE_NORM_FLOOR = 0.25 + _EPS = 1.0e-7 + _STAT_EPS = 1.0e-12 + # Frames drawn per sampled system for the calibration. The count trades + # start-up time, linear in it, against the spread of a sample mean. On a + # variable-size store the available frames follow the training batch-size + # specification and ``data_stat_nbatch``, which for OMat24 leave the pool + # far larger than this count. + _STAT_FRAMES_PER_SAMPLE = 64 + _COMPRESSION_BUFFER_NAMES = ( + "data", + "info", + "pair_film", + "pair_mixing", + "type_embedding", + "readout_matrices", + "coupling_meta", + "coupling_entry", + "coupling_value", + "output_mean", + "output_inv_std", + ) + + def __init__( + self, + rcut: float, + ntypes: int, + channels: int = 32, + lmax: int = 2, + basis_type: str = "bessel", + n_radial: int = 16, + radial_modes: int = 0, + use_amp: bool = False, + exclude_types: list[tuple[int, int]] = [], + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + spin: None = None, + ) -> None: + # === Step 1. Validate the public architecture contract === + if spin is not None: + raise NotImplementedError("DPA4C does not support spin inputs.") + if rcut <= 0.0: + raise ValueError(f"`rcut` must be positive, got {rcut}") + if ntypes <= 0: + raise ValueError(f"`ntypes` must be positive, got {ntypes}") + if n_radial <= 0: + raise ValueError(f"`n_radial` must be positive, got {n_radial}") + if ( + not isinstance(radial_modes, int) + or isinstance(radial_modes, bool) + or radial_modes < 0 + ): + raise ValueError("`radial_modes` must be a non-negative integer.") + # `channels` and `lmax` are validated inside the profile derivation, + # which owns their supported sets. + degree_channels = derive_degree_channels(channels, lmax) + bispectrum_ranks = derive_bispectrum_ranks(degree_channels) + + # === Step 2. Resolve the scalar configuration === + self.rcut = float(rcut) + self.ntypes = int(ntypes) + self.channels = int(channels) + self.lmax = int(lmax) + self.degree_channels = degree_channels + self.bispectrum_ranks = bispectrum_ranks + self.basis_type = str(basis_type).lower() + self.n_radial = int(n_radial) + self.radial_modes = int(radial_modes) + self.use_amp = bool(use_amp) + self.precision = precision + self.trainable = bool(trainable) + self.type_map = type_map + self.seed = seed + radial_hidden = resolve_swiglu_hidden_width(self.channels) + + # === Step 3. Build the shared DPA4 edge representation === + # The radial basis is raw. One p=5 DPA4 envelope gates the complete + # role-conditioned amplitude, preserving a single C³ cutoff factor. + # + # Child seeds are numbered contiguously in construction order, so each + # trainable component owns one slot and no two components can collide. + self.type_embedding = SeZMTypeEmbedding( + ntypes=self.ntypes, + embed_dim=self.channels, + precision=self.precision, + seed=child_seed(seed, 0), + trainable=self.trainable, + padding=True, + ) + self.radial_basis = RadialBasis( + rcut=self.rcut, + basis_type=self.basis_type, + n_radial=self.n_radial, + precision=self.precision, + exponent=self._ENVELOPE_EXPONENT, + apply_envelope=False, + ) + self.radial_embedding = SwiGLUMLP( + [self.n_radial, radial_hidden, self.channels], + precision=self.precision, + trainable=self.trainable, + seed=child_seed(seed, 1), + ) + # The mode profiles branch off the shared radial hidden state, so the + # residual costs one linear head rather than a second trunk. + self.radial_mode_head = ( + NativeLayer( + radial_hidden, + self.radial_modes, + bias=False, + precision=self.precision, + trainable=self.trainable, + seed=child_seed(seed, 2), + ) + if self.radial_modes > 0 + else None + ) + self.edge_envelope = C3CutoffEnvelope( + rcut=self.rcut, + exponent=self._ENVELOPE_EXPONENT, + precision=self.precision, + ) + self.pair_film = OrderedPairFiLM( + self.channels, + radial_modes=self.radial_modes, + precision=self.precision, + trainable=self.trainable, + seed=child_seed(seed, 3), + ) + self.readout = InvariantReadout( + self.channels, + self.lmax, + precision=self.precision, + trainable=self.trainable, + seed=child_seed(seed, 4), + ) + + # === Step 4. Lay out the flat moment payload === + # Degree zero owns the leading `channels` entries of the flat layout, + # so the non-scalar block is exactly its complement and needs no + # separate degree index. + channel_index, harmonic_index = build_moment_indices(self.degree_channels) + self.angular_channel_index = channel_index[self.channels :] + self.angular_harmonic_index = harmonic_index[self.channels :] + + # === Step 5. Initialize the output calibration state === + mean = np.zeros( + self.get_dim_out(), + dtype=PRECISION_DICT[self.precision], + ) + self.mean = mean + self.stddev = np.ones_like(mean) + self.compress = False + self.reinit_exclude(exclude_types) + + # === Descriptor evaluation === + + def call_graph( + self, + graph: Any, + atype: Array, + type_embedding: Array | None = None, + comm_dict: dict | None = None, + ) -> tuple[Array, None]: + """Evaluate DPA4C on a flat neighbor graph. + + Parameters + ---------- + graph + Neighbor graph containing ``edge_index`` with shape ``(2, E)``, + ``edge_vec`` with shape ``(E, 3)`` in Å, and ``edge_mask`` with + shape ``(E,)``. + atype + Flat node types with shape ``(N,)``. Padding nodes use type index + ``ntypes`` and therefore gather the zero type-embedding row. + type_embedding + Optional precomputed DPA4 type table with shape + ``(ntypes + 1, channels)``. If omitted, the descriptor + evaluates its type-embedding module. + comm_dict + Communication metadata accepted by the common graph ABI. DPA4C + does not read source-node features, so no halo-feature exchange is + required and this argument is unused. + + Returns + ------- + descriptor + Rotation- and permutation-invariant node features with shape + ``(N, get_dim_out())`` and the same floating dtype as ``edge_vec``. + rot_mat + ``None``. DPA4C does not expose an equivariant fitting input. + """ + del comm_dict + # === Step 1. Resolve type features and compute precision === + if type_embedding is None: + type_embedding = self.type_embedding.call() + xp = array_api_compat.array_namespace(graph.edge_vec) + in_dtype = graph.edge_vec.dtype + compute_dtype = get_xp_precision(xp, self.precision) + if in_dtype != compute_dtype: + graph = dataclasses.replace( + graph, + edge_vec=xp.astype(graph.edge_vec, compute_dtype), + ) + + # === Step 2. Evaluate the graph-native equations === + descriptor, _ = self.evaluate_graph(graph, atype, type_embedding) + + # === Step 3. Restore the graph input dtype === + if descriptor.dtype != in_dtype: + descriptor = xp.astype(descriptor, in_dtype) + return descriptor, None + + @cast_precision + def call( + self, + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + comm_dict: dict | None = None, + charge_spin: Array | None = None, + ) -> tuple[Array, None, None, None, Array]: + """Adapt a bounded dense neighbor list to the graph-native equations. + + This method exists for the common descriptor ABI and numerical + reference tests. Production DPA4C execution uses :meth:`call_graph` + with a carry-all graph. A rectangular list at the internal compatibility + capacity is rejected because its completeness cannot be established. + + Parameters + ---------- + coord_ext + Extended coordinates with shape ``(F, N_all, 3)`` or + ``(F, 3 * N_all)`` in Å. + atype_ext + Extended atom types with shape ``(F, N_all)``. + nlist + Bounded neighbor list with shape ``(F, N_local, N_slot)``. Negative + indices denote padding. + mapping + Extended-to-local owner mapping with shape ``(F, N_all)``. ``None`` + denotes the identity mapping. + fparam + Frame parameters accepted by the common descriptor ABI; unused. + comm_dict + Communication metadata accepted by the common descriptor ABI; + unused. + charge_spin + Charge/spin conditioning accepted by the common descriptor ABI; + unsupported and unused. + + Returns + ------- + descriptor + Invariant features with shape + ``(F, N_local, get_dim_out())``. + rot_mat + ``None``. + g2 + ``None``. + h2 + ``None``. + envelope + Per-slot C³ envelope with shape + ``(F, N_local, N_slot, 1)``. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + graph_from_dense_quartet, + ) + + del fparam, comm_dict, charge_spin + xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) + nf, nloc, nnei = nlist.shape + + # === Step 1. Convert the dense quartet without compacting its edge axis === + graph, atype_local = graph_from_dense_quartet( + coord_ext, + atype_ext, + nlist, + mapping, + ) + + # === Step 2. Evaluate the same graph-native equations === + descriptor, envelope = self.evaluate_graph( + graph, + atype_local, + self.type_embedding.call(), + ) + + # === Step 3. Restore the common dense descriptor ABI === + descriptor = xp.reshape( + descriptor, + (nf, nloc, descriptor.shape[-1]), + ) + envelope = xp.reshape(envelope, (nf, nloc, nnei, 1)) + return descriptor, None, None, None, envelope + + def evaluate_graph( + self, + graph: Any, + atype: Array, + type_embedding: Array, + ) -> tuple[Array, Array]: + """Evaluate the graph-native descriptor equations. + + The pipeline is edge amplitudes, one destination reduction into + degree-wise moments, and the fixed invariant readout. + + Parameters + ---------- + graph + Neighbor graph in descriptor compute precision. + atype + Flat node types with shape ``(N,)``. + type_embedding + Complete type table with shape ``(ntypes + 1, channels)``. + + Returns + ------- + descriptor + Invariant node features with shape ``(N, get_dim_out())``. + envelope + Masked per-edge C³ envelope with shape ``(E, 1)``. + """ + xp = array_api_compat.array_namespace(graph.edge_vec) + + # === Step 1. Place the precomputed type table in the graph namespace === + # A converted dpmodel may already store the table in the active + # namespace. Conversion is required only for a direct NumPy-defined + # descriptor evaluated with JAX or another array backend. + type_namespace = array_api_compat.array_namespace(type_embedding) + if type_namespace is not xp: + type_embedding = xp.asarray( + type_embedding, + dtype=graph.edge_vec.dtype, + device=array_api_compat.device(graph.edge_vec), + ) + dst = graph.edge_index[1] + n_total = atype.shape[0] + center_type_embedding = self.gather_rows(type_embedding, atype, xp) + pair_tables = self.pair_film.call(type_embedding) + + # === Step 2. Build the masked edge amplitudes and harmonics === + amplitude, basis, envelope = self.build_edge_features( + graph, + atype, + *pair_tables, + ) + + # === Step 3. Reduce the degree-wise moments === + moments, divisors = self.aggregate_moments( + amplitude, + basis, + envelope, + dst, + n_total, + ) + + # === Step 4. Build calibrated invariant features === + return ( + self.build_invariant_descriptor( + moments, + center_type_embedding, + divisors, + ), + envelope[:, None], + ) + + def build_edge_features( + self, + graph: Any, + atype: Array, + pair_scale: Array, + pair_shift: Array, + pair_mixing: Array | None, + ) -> tuple[Array, Array, Array]: + r"""Build the enveloped edge amplitudes and the masked harmonics. + + The ordered type pair :math:`(a,b)` rescales the one shared radial + function :math:`g` and mixes the :math:`R` shared mode profiles + :math:`q_\mu` with its own coefficients: + + .. math:: + + \phi_{ijc} + =\chi_{ij}\Bigl( + \gamma_{ab,c}g_c(\rho_{ij})+\beta_{ab,c} + +\sum_\mu U_{ab,c\mu}q_\mu(\rho_{ij}) + \Bigr). + + Parameters + ---------- + graph + Neighbor graph in descriptor compute precision. + atype + Flat node types with shape ``(N,)``. + pair_scale + Ordered radial scales with shape + ``((ntypes + 1) ** 2, channels)``. + pair_shift + Ordered radial shifts with shape + ``((ntypes + 1) ** 2, channels)``. + pair_mixing + Ordered mode-mixing table with shape + ``((ntypes + 1) ** 2, channels, radial_modes)``, or ``None`` when + ``radial_modes`` is zero. + + Returns + ------- + amplitude + Masked edge amplitudes with shape ``(E, channels)``. + basis + Masked Cartesian harmonics with shape ``(E, (lmax + 1) ** 2)``. + envelope + Masked C³ envelope with shape ``(E,)``. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, + ) + + # === Step 1. Merge graph and descriptor-level exclusion masks === + graph = apply_pair_exclusion(graph, atype, self.emask) + xp = array_api_compat.array_namespace(graph.edge_vec) + src, dst = graph.edge_index[0], graph.edge_index[1] + center_type = self.gather_rows(atype, dst, xp) + neighbor_type = self.gather_rows(atype, src, xp) + + # === Step 2. Build regularized distances and directions === + # sqrt(r^2 + eps^2) keeps the direction finite for coincident or guard + # edges. Valid physical edges are unaffected above the 1e-7 Å scale. + distance_squared = xp.sum( + graph.edge_vec * graph.edge_vec, + axis=-1, + keepdims=True, + ) + distance = xp.sqrt(distance_squared + self._EPS * self._EPS) + direction = graph.edge_vec / distance + real_type = (center_type < self.ntypes) & (neighbor_type < self.ntypes) + edge_mask = graph.edge_mask & real_type + mask = xp.astype(edge_mask[:, None], graph.edge_vec.dtype) + + # === Step 3. Evaluate the shared DPA4 radial representation === + # DPA4C requests the raw radial basis. One explicit C³ envelope gates + # the combined radial and type feature, so the scalar edge amplitude + # contains exactly one cutoff factor and vanishes smoothly at rcut. + envelope = self.evaluate_cutoff_envelope(distance) * mask + radial_basis = self.radial_basis.call(distance) + radial_hidden = self.radial_embedding.call_hidden(radial_basis) + radial = self.radial_embedding.call_output(radial_hidden) + pair_index = center_type * (self.ntypes + 1) + neighbor_type + scale = self.gather_rows(pair_scale, pair_index, xp) # (E, C) + shift = self.gather_rows(pair_shift, pair_index, xp) # (E, C) + amplitude = radial * scale + shift + + # === Step 4. Add the pair-conditioned radial mode residual === + # Each ordered pair selects its own combination of the R shared mode + # profiles. The contraction is written as a broadcast product reduced + # over the mode axis rather than as a batched matrix-vector product: + # one GEMV per edge leaves the tiny C-by-R operands far short of + # memory bandwidth, whereas the reduction is a plain streaming pass. + # Expanding the ordered table per edge dominates the cost either way. + if pair_mixing is not None: + mixing = self.gather_rows(pair_mixing, pair_index, xp) # (E, C, R) + modes = self.radial_mode_head(radial_hidden) # (E, R) + amplitude = amplitude + xp.sum(mixing * modes[:, None, :], axis=-1) + + # === Step 5. Gate the amplitude and build the masked harmonics === + return ( + amplitude * envelope, + self.build_angular_basis(direction) * mask, + envelope[:, 0], + ) + + def aggregate_moments( + self, + amplitude: Array, + basis: Array, + envelope: Array, + dst: Array, + n_total: int, + ) -> tuple[Array, Array]: + r"""Aggregate every degree-wise moment in one segment reduction. + + Degree zero is additive under the single envelope already carried by + the amplitude, whereas every non-scalar degree carries a second + envelope factor and its own matched normalizer: + + .. math:: + + d^{(0)}_i=\sum_j\chi_{ij}^2,\qquad + d^{(+)}_i=\sum_j\chi_{ij}^4,\qquad + n^{(\bullet)}_i=\bigl(d^{(\bullet)}_i+\tfrac14\bigr)^{-1/2},\\ + X^{(0)}_{ic}=n^{(0)}_i\sum_j\phi_{ijc},\qquad + X^{(\ell)}_{icm} + =n^{(+)}_i\sum_j\chi_{ij}\phi_{ijc}B^{(\ell)}_m(\hat u_{ij}). + + Both envelope masses and both moment blocks share one payload, so the + descriptor reduces the edge axis exactly once. + + Parameters + ---------- + amplitude + Masked edge amplitudes with shape ``(E, channels)``. + basis + Masked Cartesian harmonics with shape ``(E, (lmax + 1) ** 2)``. + envelope + Masked C³ envelope with shape ``(E,)``. + dst + Destination node indices with shape ``(E,)``. + n_total + Number of output nodes ``N``. + + Returns + ------- + moments + Flat normalized moments with shape ``(N, S)``, where + ``S = sum((2 * l + 1) * degree_channels[l])``. + divisors + The two divisors :math:`1/n^{(0)}` and :math:`1/n^{(+)}` with shape + ``(N, 2)``. They are retained because normalization is otherwise + irreversible: the readout sees only scaled moments and can neither + recover the unnormalized ones nor read the effective coordination + they encode. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_sum, + ) + + xp = array_api_compat.array_namespace(amplitude) + device = array_api_compat.device(amplitude) + channel_index = xp_asarray_nodetach( + xp, + self.angular_channel_index, + device=device, + ) + harmonic_index = xp_asarray_nodetach( + xp, + self.angular_harmonic_index, + device=device, + ) + + # Payload layout: [chi^2, chi^4, degree zero (C_0), degrees one and + # above (S - C_0)]. Degree zero is the whole edge amplitude under its + # single envelope, while the non-scalar block gathers an amplitude and + # a harmonic per flat moment coordinate and carries a second envelope. + envelope_squared = envelope * envelope + payload = xp.concat( + [ + envelope_squared[:, None], + (envelope_squared * envelope_squared)[:, None], + amplitude, + xp.take(amplitude, channel_index, axis=1) + * xp.take(basis, harmonic_index, axis=1) + * envelope[:, None], + ], + axis=1, + ) + reduced = segment_sum(payload, dst, n_total) + + scalar_end = 2 + self.channels + floor = self._DEGREE_NORM_FLOOR + divisors = xp.sqrt(reduced[:, :2] + floor) + return ( + xp.concat( + [ + reduced[:, 2:scalar_end] / divisors[:, :1], + reduced[:, scalar_end:] / divisors[:, 1:], + ], + axis=1, + ), + divisors, + ) + + def build_invariant_descriptor( + self, + moments: Array, + center_type_embedding: Array, + divisors: Array, + ) -> Array: + """Build and calibrate the geometric and center-type output blocks. + + The calibration is the fixed diagonal preconditioner established by + :meth:`compute_input_stats`, not a running normalization. + + Parameters + ---------- + moments + Flat degree-wise moments with shape ``(N, S)``. + center_type_embedding + Center type embeddings with shape ``(N, channels)``. + divisors + The two moment divisors with shape ``(N, 2)``. They close the + geometric block, so the calibration treats them like any other + invariant and the center-type tail keeps its trailing position. + + Returns + ------- + Array + Descriptor with shape ``(N, get_dim_out())``. + """ + xp = array_api_compat.array_namespace(moments) + device = array_api_compat.device(moments) + descriptor = xp.concat( + [self.readout.call(moments), divisors, center_type_embedding], + axis=-1, + ) + mean = xp_asarray_nodetach(xp, self.mean, device=device) + stddev = xp_asarray_nodetach(xp, self.stddev, device=device) + return (descriptor - mean[None, :]) / stddev[None, :] + + # === Backend primitives === + # A backend wrapper overrides these to reach native kernels; the equations + # above stay array-API neutral. + + def gather_rows( + self, + values: Array, + index: Array, + xp: Any | None = None, + ) -> Array: + """Gather rows along the leading axis. + + Parameters + ---------- + values + Source array with shape ``(N, ...)``. + index + Row indices with arbitrary shape. + xp + Optional array namespace for ``values``. + + Returns + ------- + Array + Gathered values with shape ``index.shape + values.shape[1:]``. + """ + if xp is None: + xp = array_api_compat.array_namespace(values) + return xp.take(values, index, axis=0) + + def evaluate_cutoff_envelope(self, distance: Array) -> Array: + """Evaluate the fixed C³ cutoff envelope. + + Parameters + ---------- + distance + Regularized edge distances with shape ``(E, 1)`` in Å. + + Returns + ------- + Array + Envelope values with shape ``(E, 1)``. + """ + return self.edge_envelope.call(distance) + + def build_angular_basis(self, direction: Array) -> Array: + """Build real Cartesian harmonics through ``lmax``. + + Parameters + ---------- + direction + Regularized edge directions with shape ``(E, 3)``. + + Returns + ------- + Array + Packed harmonics with shape ``(E, (lmax + 1) ** 2)``. + """ + return build_angular_basis(direction, self.lmax) + + # === Parameter sharing and statistics === + + def share_params( + self, + base_class: Any, + shared_level: int, + model_prob: float = 1.0, + resume: bool = False, + ) -> None: + """Share all descriptor parameters for multitask training. + + Parameters + ---------- + base_class + DPA4C descriptor that owns the shared parameters. + shared_level + Sharing level. DPA4C supports only level ``0``, which shares the + complete descriptor. + model_prob + Model sampling probability accepted by the common multitask ABI; + unused because DPA4C has no mergeable input statistics. + resume + Whether sharing occurs during checkpoint restoration; unused. + """ + del model_prob, resume + if self.__class__ != base_class.__class__: + raise TypeError("Only DPA4C descriptors can share parameters.") + signature = self.structure_signature() + base_signature = base_class.structure_signature() + if signature != base_signature: + raise ValueError( + "DPA4C parameter sharing requires identical structural " + f"parameters, got {signature} and {base_signature}" + ) + if shared_level != 0: + raise NotImplementedError("DPA4C supports only shared_level=0.") + for name in ( + "type_embedding", + "radial_basis", + "radial_embedding", + "radial_mode_head", + "pair_film", + "readout", + ): + setattr(self, name, getattr(base_class, name)) + self.mean = base_class.mean + self.stddev = base_class.stddev + + def structure_signature(self) -> tuple: + """Return the configuration that must agree between sharing replicas. + + The signature compares two live descriptors and is never persisted, + so it may reference execution policy as well as persisted structure. + + It covers every field that fixes the shape or the meaning + of a shared module. ``rcut``, ``basis_type``, and ``n_radial`` define + the radial basis; ``ntypes`` defines the type table and the ordered + pair index space; ``channels``, ``lmax``, and ``radial_modes`` define + every remaining width; ``use_amp`` selects the precision policy that a + backend attaches to the shared layers, so a replica that autocasts + against layers configured without it would silently lose the effect; + ``trainable`` decides whether + those layers carry gradients at all; ``type_map`` fixes what the rows + of the shared type table mean. Precision itself enters through its + resolved dtype so that equivalent spellings agree. + + Branch-local state is deliberately absent. ``exclude_types`` is the + only such field: it configures the pair-exclusion mask, which each + replica keeps for itself. + + Returns + ------- + tuple + Structural configuration of the shareable modules. + """ + return ( + self.rcut, + self.ntypes, + self.channels, + self.lmax, + self.basis_type, + self.n_radial, + self.radial_modes, + self.use_amp, + self.trainable, + None if self.type_map is None else tuple(self.type_map), + np.dtype(PRECISION_DICT[self.precision]).name, + ) + + def change_type_map( + self, + type_map: list[str], + model_with_new_type_stat: Any | None = None, + ) -> None: + """Reject unsupported atom-type remapping. + + Parameters + ---------- + type_map + Requested atom-type map. + model_with_new_type_stat + Optional descriptor carrying statistics for newly introduced + types. DPA4C does not use descriptor input statistics. + """ + del type_map, model_with_new_type_stat + raise NotImplementedError("DPA4C does not support changing `type_map`.") + + def set_stat_mean_and_stddev(self, mean: Array, stddev: Array) -> None: + """Store fixed output calibration arrays. + + Parameters + ---------- + mean + Output shift with shape ``(get_dim_out(),)``. + stddev + Positive output scale with shape ``(get_dim_out(),)``. + """ + expected = (self.get_dim_out(),) + if mean.shape != expected or stddev.shape != expected: + raise ValueError( + "DPA4C output statistics must both have shape " + f"{expected}, got {mean.shape} and {stddev.shape}" + ) + if np.any(to_numpy_array(stddev) <= 0.0): + raise ValueError("DPA4C output scales must be positive.") + self.mean = mean + self.stddev = stddev + + def get_stat_mean_and_stddev(self) -> tuple[Array, Array]: + """Return interface-compatible descriptor statistics. + + Returns + ------- + mean + Stored mean array. + stddev + Stored standard-deviation array. + """ + return self.mean, self.stddev + + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, + ) -> None: + """Calibrate polynomial output families against the type-embedding RMS. + + The calibration is a fixed initialization preconditioner. Every + geometric coordinate is measured independently, while the center type + tail is left unchanged. No sample-dependent normalization is evaluated + during training or inference. + + Parameters + ---------- + merged + Sampled training systems or a callable returning them. + path + Optional statistics path. Model-dependent calibration is always + recomputed and therefore does not consume this path. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + del path + sampled = merged() if callable(merged) else merged + if not sampled: + return + + xp = array_api_compat.array_namespace(self.stddev) + device = array_api_compat.device(self.stddev) + dtype = self.stddev.dtype + mean_backup, stddev_backup = self.mean, self.stddev + self.mean = xp.zeros_like(self.mean) + self.stddev = xp.ones_like(self.stddev) + geometry_dim = self.get_dim_out() - self.channels + square_sum = np.zeros(geometry_dim, dtype=np.float64) + value_sum = np.zeros(geometry_dim, dtype=np.float64) + value_count = 0 + + try: + for system in sampled: + coord_np = to_numpy_array(system["coord"]) + nframes = coord_np.shape[0] + coord_np = np.reshape(coord_np, (nframes, -1, 3)) + atype_np = np.reshape( + to_numpy_array(system["atype"]), + (nframes, -1), + ) + box_value = system.get("box", None) + box_np = ( + None + if box_value is None + else np.reshape(to_numpy_array(box_value), (nframes, -1)) + ) + nstat_frames = min(nframes, self._STAT_FRAMES_PER_SAMPLE) + frame_indices = np.linspace( + 0, + nframes - 1, + num=nstat_frames, + dtype=np.int64, + ) + for frame_index in frame_indices: + coord = xp.asarray( + coord_np[frame_index : frame_index + 1], + dtype=dtype, + device=device, + ) + atype = xp.asarray( + atype_np[frame_index : frame_index + 1], + device=device, + ) + box = ( + None + if box_np is None + else xp.asarray( + box_np[frame_index : frame_index + 1], + dtype=dtype, + device=device, + ) + ) + graph = build_neighbor_graph( + coord, + atype, + box, + self.get_rcut(), + ) + output, _ = self.call_graph(graph, xp.reshape(atype, (-1,))) + output_np = to_numpy_array(output).reshape( + -1, + self.get_dim_out(), + ) + if output_np.shape[0] == 0: + continue + square_sum += np.sum( + np.square( + output_np[:, :geometry_dim], + dtype=np.float64, + ), + axis=0, + dtype=np.float64, + ) + value_sum += np.sum( + output_np[:, :geometry_dim], + axis=0, + dtype=np.float64, + ) + value_count += output_np.shape[0] + finally: + self.mean, self.stddev = mean_backup, stddev_backup + + if value_count == 0: + return + feature_rms = np.sqrt(square_sum / float(value_count)) + if np.any(~np.isfinite(feature_rms)) or np.any(feature_rms <= self._STAT_EPS): + raise ValueError( + "DPA4C output calibration requires non-degenerate finite " + f"features, got RMS values {feature_rms.tolist()}" + ) + type_table = to_numpy_array(self.type_embedding.call())[: self.ntypes] + target_rms = float(np.sqrt(np.mean(np.square(type_table, dtype=np.float64)))) + if not math.isfinite(target_rms) or target_rms <= self._STAT_EPS: + raise ValueError( + f"DPA4C type embedding has a degenerate calibration RMS {target_rms}" + ) + geometry_stddev = feature_rms / target_rms + geometry_mean = np.zeros(geometry_dim, dtype=np.float64) + + # The two moment divisors are the only outputs carrying their + # information on a large offset rather than around zero: their RMS + # exceeds a typical invariant by two orders of magnitude, so scaling + # alone would leave them at one plus a small fluctuation. They are + # standardized; every other coordinate keeps the shared RMS + # preconditioner, whose zero mean the readout construction justifies. + mass = slice(geometry_dim - 2, geometry_dim) + mass_mean = value_sum[mass] / float(value_count) + mass_stddev = np.sqrt( + np.maximum( + square_sum[mass] / float(value_count) - np.square(mass_mean), 0.0 + ) + ) + if np.any(mass_stddev <= self._STAT_EPS): + raise ValueError( + "DPA4C neighborhood masses are constant over the calibration " + f"sample, got standard deviations {mass_stddev.tolist()}" + ) + geometry_mean[mass] = mass_mean + geometry_stddev[mass] = mass_stddev / target_rms + + tail = np.zeros(self.channels, dtype=np.float64) + self.mean = np.concatenate([geometry_mean, tail]).astype( + PRECISION_DICT[self.precision] + ) + self.stddev = np.concatenate([geometry_stddev, tail + 1.0]).astype( + PRECISION_DICT[self.precision] + ) + + # === Serialization and neighbor statistics === + + def serialize(self) -> dict: + """Serialize the descriptor. + + ``use_amp`` is deliberately absent. It selects an execution policy + rather than any part of the learned function, so it is supplied by the + training configuration and by ``DP_AMP_INFER``, never restored from a + checkpoint. + + Returns + ------- + dict + Versioned descriptor configuration, nested DPA4 components, and + interface statistics. + """ + data = { + "@class": "Descriptor", + "type": "dpa4c", + "@version": 1, + "rcut": self.rcut, + "ntypes": self.ntypes, + "channels": self.channels, + "lmax": self.lmax, + "basis_type": self.basis_type, + "n_radial": self.n_radial, + "radial_modes": self.radial_modes, + "exclude_types": self.exclude_types, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, + "type_map": self.type_map, + "seed": self.seed, + "spin": None, + "type_embedding": self.type_embedding.serialize(), + "radial_basis": self.radial_basis.serialize(), + "radial_embedding": self.radial_embedding.serialize(), + "radial_mode_head": ( + None + if self.radial_mode_head is None + else self.radial_mode_head.serialize() + ), + "pair_film": self.pair_film.serialize(), + "readout": self.readout.serialize(), + "@variables": { + "mean": to_numpy_array(self.mean), + "stddev": to_numpy_array(self.stddev), + }, + } + if self.compress: + data["compress"] = { + "@variables": { + name: to_numpy_array(getattr(self, f"compress_{name}")) + for name in self._COMPRESSION_BUFFER_NAMES + } + } + return data + + @classmethod + def deserialize(cls, data: dict) -> DescrptDPA4C: + """Deserialize a DPA4C descriptor. + + Parameters + ---------- + data + Versioned descriptor dictionary produced by :meth:`serialize`. + + Returns + ------- + DescrptDPA4C + Reconstructed descriptor with restored trainable components. + """ + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + data.pop("type") + compression = data.pop("compress", None) + variables = data.pop("@variables") + type_embedding = data.pop("type_embedding") + radial_basis = data.pop("radial_basis") + radial_embedding = data.pop("radial_embedding") + radial_mode_head = data.pop("radial_mode_head") + pair_film = data.pop("pair_film") + readout = data.pop("readout") + + obj = cls(**data) + obj.type_embedding = SeZMTypeEmbedding.deserialize(type_embedding) + obj.radial_basis = RadialBasis.deserialize(radial_basis) + obj.radial_embedding = SwiGLUMLP.deserialize(radial_embedding) + obj.radial_mode_head = ( + None + if radial_mode_head is None + else NativeLayer.deserialize(radial_mode_head) + ) + obj.pair_film = OrderedPairFiLM.deserialize(pair_film) + obj.readout = InvariantReadout.deserialize(readout) + obj.set_stat_mean_and_stddev( + variables["mean"], + variables["stddev"], + ) + if compression is not None: + obj._load_compression(compression) + return obj + + def _load_compression(self, compression: dict) -> None: + """Restore the compressed-inference artifact payload.""" + variables = compression["@variables"] + for name in self._COMPRESSION_BUFFER_NAMES: + setattr(self, f"compress_{name}", variables[name]) + self.compress = True + + @classmethod + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: list[str] | None, + local_jdata: dict, + ) -> tuple[dict, float]: + """Report the minimum neighbor distance without introducing a ``sel``. + + The descriptor is graph-native, so no neighbor capacity has to be + derived from the data and the returned configuration is unchanged. + Only the minimum neighbor distance is measured, which downstream + consumers use to bound tabulated radial ranges. + + Parameters + ---------- + train_data + Training dataset used for neighbor statistics. + type_map + Ordered atom-type names. + local_jdata + DPA4C descriptor configuration. + + Returns + ------- + local_jdata + Unmodified descriptor configuration. + min_nbor_dist + Minimum observed neighbor distance in Å. + """ + del type_map + return local_jdata.copy(), cls._update_sel_cls().get_min_nbor_dist(train_data) + + # === Common descriptor ABI === + + @property + def dim_out(self) -> int: + """Return the invariant descriptor width.""" + return self.get_dim_out() + + def get_rcut(self) -> float: + """Return the outer cutoff radius in Å.""" + return self.rcut + + def get_rcut_smth(self) -> float: + """Return the outer cutoff used as the common smoothing radius.""" + return self.rcut + + def get_sel(self) -> list[int]: + """Return an effectively unbounded neighbor capacity. + + The descriptor is graph-native and carries every neighbor within the + cutoff, so it imposes no capacity. The common ``sel`` ABI still + requires a number; a value no environment can reach reports that + absence of a bound. + """ + return [999999] + + def get_ntypes(self) -> int: + """Return the number of real atom types.""" + return self.ntypes + + def get_type_map(self) -> list[str] | None: + """Return the ordered atom-type names.""" + return self.type_map + + def get_dim_out(self) -> int: + r"""Return the complete invariant descriptor width ``D``. + + The flat equivariant moment state has width + + .. math:: + + S=\sum_{\ell=0}^{L}(2\ell+1)C_\ell, + + where ``C_l = degree_channels[l]`` is derived from the scalar + ``channels`` parameter. ``S`` controls edge-reduction work and + moment-state memory, but these equivariant coefficients are not + exposed directly to the fitting network. + + The invariant output width is + + .. math:: + + D=2C_0 + +\sum_{\ell=1}^{L}\frac{C_\ell(C_\ell+1)}{2} + +D_{\mathrm{bispectrum}} + +K_1K_2 + +2. + + The two ``C_0`` blocks are the scalar moments and the center-type + embedding. The summation contains the exact upper-triangular Gram for + each non-scalar degree, ``K_l`` denotes the derived bispectrum rank, + the quartic term is the projected ``|Q_b v_a|^2``, and the trailing + pair is the two neighborhood masses. The mixing rank does not enter + ``D``. + + ``D_bispectrum`` is the sum over allowed O(3)-even degree triples: + + - distinct degrees contribute ``K_l1 * K_l2 * K_l3``; + - two equal degrees contribute one symmetric pair count + ``K * (K + 1) // 2`` times the remaining rank; + - three equal degrees contribute ``K * (K + 1) * (K + 2) // 6``. + + Returns + ------- + int + Width of the invariant descriptor consumed by the fitting network. + """ + return self.readout.get_dim_out() + self.channels + 2 + + def get_dim_emb(self) -> int: + """Return zero because fitting receives no equivariant channels.""" + return 0 + + def mixed_types(self) -> bool: + """Return whether the descriptor consumes a mixed-type neighbor list.""" + return True + + def has_message_passing(self) -> bool: + """Return whether source-node features are exchanged.""" + return False + + def has_message_passing_across_ranks(self) -> bool: + """Return whether intermediate halo communication is required.""" + return False + + def need_sorted_nlist_for_lower(self) -> bool: + """Return whether graph-lower edges must be destination sorted.""" + return False + + def get_env_protection(self) -> float: + """Return the direction regularization scale in Å.""" + return self._EPS + + def uses_graph_lower(self) -> bool: + """Return whether graph-native lowering is supported.""" + return True + + def graph_edge_dtype(self) -> str: + """Return the edge-geometry dtype accepted by graph deployment. + + Returns + ------- + str + ``"float32"`` for compressed float32 inference, otherwise + ``"float64"``. + """ + precision = np.dtype(PRECISION_DICT[self.precision]).name + return "float32" if self.compress and precision == "float32" else "float64" + + def reinit_exclude( + self, + exclude_types: list[tuple[int, int]] | None = None, + ) -> None: + """Rebuild the ordered pair-exclusion mask.""" + if exclude_types is None: + exclude_types = [] + self.exclude_types = list(exclude_types) + self.emask = PairExcludeMask( + self.ntypes, + exclude_types=self.exclude_types, + ) diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py new file mode 100644 index 0000000000..4d5d2fbdbf --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Neural and geometric building blocks for DPA4C.""" + +from .bispectrum import ( + BispectrumLayout, + build_bispectrum_layout, + derive_bispectrum_ranks, + enumerate_degree_triples, +) +from .geometry import ( + MAX_ANGULAR_DEGREE, + build_angular_basis, + build_moment_indices, + degree_offsets, + derive_degree_channels, + packed_l2_to_stf, +) +from .pair_film import ( + OrderedPairFiLM, +) +from .readout import ( + InvariantReadout, +) + +__all__ = [ + "MAX_ANGULAR_DEGREE", + "BispectrumLayout", + "InvariantReadout", + "OrderedPairFiLM", + "build_angular_basis", + "build_bispectrum_layout", + "build_moment_indices", + "degree_offsets", + "derive_bispectrum_ranks", + "derive_degree_channels", + "enumerate_degree_triples", + "packed_l2_to_stf", +] diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/bispectrum.py b/deepmd/dpmodel/descriptor/dpa4c_nn/bispectrum.py new file mode 100644 index 0000000000..c1885bdacf --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/bispectrum.py @@ -0,0 +1,315 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""O(3)-invariant Cartesian bispectrum coupling tables.""" + +from __future__ import ( + annotations, +) + +import functools +import itertools +import math +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import numpy as np + +from deepmd.dpmodel.utils.lebedev import ( + load_lebedev_rule, +) + +from .geometry import ( + MAX_ANGULAR_DEGREE, + build_angular_basis, +) + +if TYPE_CHECKING: + from collections.abc import ( + Sequence, + ) + +DegreeTriple = tuple[int, int, int] + + +@dataclass(frozen=True) +class BispectrumLayout: + """Store flattened coupling and independent probe-index tables.""" + + degree_triples: tuple[DegreeTriple, ...] + coupling: np.ndarray + coupling_offsets: tuple[int, ...] + probe_index: np.ndarray + probe_scale: np.ndarray + probe_offsets: tuple[int, ...] + + @property + def dim_out(self) -> int: + """Return the number of independent probe contractions. + + Returns + ------- + int + Total bispectrum feature count. + """ + return int(self.probe_index.size) + + +def enumerate_degree_triples(lmax: int) -> tuple[DegreeTriple, ...]: + """Enumerate non-scalar O(3)-even bispectrum degree triples. + + Parameters + ---------- + lmax + Maximum angular degree. Supported values are zero through four. + + Returns + ------- + tuple[DegreeTriple, ...] + Sorted triples satisfying the triangle and even-parity conditions. + + Raises + ------ + ValueError + If ``lmax`` is outside the supported range. + """ + if lmax < 0 or lmax > MAX_ANGULAR_DEGREE: + raise ValueError( + f"`lmax` must be between 0 and {MAX_ANGULAR_DEGREE}, got {lmax}" + ) + triples = [] + for degree_1 in range(1, lmax + 1): + for degree_2 in range(degree_1, lmax + 1): + for degree_3 in range(degree_2, lmax + 1): + if degree_3 > degree_1 + degree_2: + continue + if (degree_1 + degree_2 + degree_3) % 2 != 0: + continue + triples.append((degree_1, degree_2, degree_3)) + return tuple(triples) + + +@functools.lru_cache(maxsize=5) +def _coupling_tables( + lmax: int, +) -> tuple[tuple[DegreeTriple, ...], tuple[np.ndarray, ...]]: + """Build unit-norm real Cartesian Gaunt tensors.""" + triples = enumerate_degree_triples(lmax) + if not triples: + return triples, () + + points, weights = load_lebedev_rule(13) + basis = np.asarray(build_angular_basis(points, lmax), dtype=np.float64) + couplings = [] + for degree_1, degree_2, degree_3 in triples: + block_1 = basis[:, degree_1**2 : (degree_1 + 1) ** 2] + block_2 = basis[:, degree_2**2 : (degree_2 + 1) ** 2] + block_3 = basis[:, degree_3**2 : (degree_3 + 1) ** 2] + coupling = np.einsum( + "n,ni,nj,nk->ijk", + weights, + block_1, + block_2, + block_3, + optimize=True, + ) + coupling = _symmetrize_equal_degrees( + coupling, + (degree_1, degree_2, degree_3), + ) + norm = float(np.linalg.norm(coupling)) + if not math.isfinite(norm) or norm <= 1.0e-14: + raise ValueError( + "Degenerate bispectrum coupling for degree triple " + f"{(degree_1, degree_2, degree_3)}" + ) + coupling /= norm + significant = np.flatnonzero(np.abs(coupling) > 1.0e-14) + if significant.size > 0 and coupling.flat[significant[0]] < 0.0: + coupling = -coupling + couplings.append(coupling) + return triples, tuple(couplings) + + +def _symmetrize_equal_degrees( + coupling: np.ndarray, + degrees: DegreeTriple, +) -> np.ndarray: + """Enforce permutation symmetry for equal-degree tensor axes.""" + if degrees[0] == degrees[2]: + permutations = tuple(itertools.permutations((0, 1, 2))) + elif degrees[0] == degrees[1]: + permutations = ((0, 1, 2), (1, 0, 2)) + elif degrees[1] == degrees[2]: + permutations = ((0, 1, 2), (0, 2, 1)) + else: + return coupling + return sum(np.transpose(coupling, axes) for axes in permutations) / len( + permutations + ) + + +def _probe_entries( + degrees: DegreeTriple, + ranks: Sequence[int], +) -> tuple[np.ndarray, np.ndarray]: + """Build independent flattened probe indices and isometric scales. + + A symmetrized coupling tensor makes the contraction invariant under any + permutation of the axes that carry equal degrees, so only one + representative of each orbit is emitted. Scaling that representative by + the square root of its orbit size keeps the reduced feature vector + isometric to the full one. + """ + entries: list[tuple[int, int, int]] = [] + scales: list[float] = [] + rank_1, rank_2, rank_3 = (int(ranks[degree - 1]) for degree in degrees) + + if degrees[0] == degrees[2]: + for entry in itertools.combinations_with_replacement( + range(rank_1), + 3, + ): + entries.append(entry) + counts = [entry.count(value) for value in set(entry)] + multiplicity = math.factorial(3) + for count in counts: + multiplicity //= math.factorial(count) + scales.append(math.sqrt(float(multiplicity))) + elif degrees[0] == degrees[1]: + for first, second in itertools.combinations_with_replacement( + range(rank_1), + 2, + ): + scale = 1.0 if first == second else math.sqrt(2.0) + for third in range(rank_3): + entries.append((first, second, third)) + scales.append(scale) + elif degrees[1] == degrees[2]: + for first in range(rank_1): + for second, third in itertools.combinations_with_replacement( + range(rank_2), + 2, + ): + entries.append((first, second, third)) + scales.append(1.0 if second == third else math.sqrt(2.0)) + else: + entries.extend( + itertools.product( + range(rank_1), + range(rank_2), + range(rank_3), + ) + ) + scales.extend([1.0] * (rank_1 * rank_2 * rank_3)) + + index = np.asarray( + [ + (first * rank_2 + second) * rank_3 + third + for first, second, third in entries + ], + dtype=np.int64, + ) + return index, np.asarray(scales, dtype=np.float64) + + +def build_bispectrum_layout( + lmax: int, + ranks: Sequence[int], +) -> BispectrumLayout: + """Build flattened coupling and independent probe-output tables. + + Parameters + ---------- + lmax + Maximum angular degree. Supported values are zero through four. + ranks + Effective positive probe widths for degrees one through ``lmax``. + + Returns + ------- + BispectrumLayout + Immutable layout arrays used by all array backends. + + Raises + ------ + ValueError + If the rank profile does not match ``lmax``. + """ + triples = enumerate_degree_triples(lmax) + ranks = tuple(int(rank) for rank in ranks) + if len(ranks) != lmax: + raise ValueError(f"`ranks` must contain {lmax} entries, got {len(ranks)}") + if any(rank <= 0 for rank in ranks): + raise ValueError(f"`ranks` must be positive, got {ranks}") + if not triples: + return BispectrumLayout( + degree_triples=(), + coupling=np.empty(0, dtype=np.float64), + coupling_offsets=(0,), + probe_index=np.empty(0, dtype=np.int64), + probe_scale=np.empty(0, dtype=np.float64), + probe_offsets=(0,), + ) + + _, coupling_tables = _coupling_tables(lmax) + coupling_offsets = [0] + probe_offsets = [0] + coupling_parts = [] + probe_index_parts = [] + probe_scale_parts = [] + for degrees, coupling in zip(triples, coupling_tables, strict=True): + coupling_parts.append(np.reshape(coupling, (-1,))) + coupling_offsets.append(coupling_offsets[-1] + coupling.size) + probe_index, probe_scale = _probe_entries(degrees, ranks) + probe_index_parts.append(probe_index) + probe_scale_parts.append(probe_scale) + probe_offsets.append(probe_offsets[-1] + probe_index.size) + + return BispectrumLayout( + degree_triples=triples, + coupling=np.concatenate(coupling_parts), + coupling_offsets=tuple(coupling_offsets), + probe_index=np.concatenate(probe_index_parts), + probe_scale=np.concatenate(probe_scale_parts), + probe_offsets=tuple(probe_offsets), + ) + + +def derive_bispectrum_ranks( + degree_channels: Sequence[int], +) -> list[int]: + """Derive the fixed degree-wise bispectrum probe ranks. + + The exact degree Gram determines a degree-one block up to the physical + rotation group, but for degree two it determines the packed coefficients + only up to O(5), of which the physical rotations form a three-parameter + subgroup. The cubic contractions resolve the remaining orientation for the + probed channels alone, so ``K_2`` trades that resolution against the width + of the cubic and quartic blocks. Raising ``K_2`` to ``C_2`` was measured to + give a small accuracy gain that does not justify the wider invariant + output, so the compact profile is fixed. + + Parameters + ---------- + degree_channels + Channel widths for degrees zero through ``lmax``. + + Returns + ------- + list[int] + Probe ranks for degrees one through ``lmax``. + + Raises + ------ + ValueError + If the degree profile does not cover a supported ``lmax``. + """ + if len(degree_channels) not in {3, 4, 5}: + raise ValueError( + "`degree_channels` must contain three to five entries, got " + f"{len(degree_channels)}" + ) + return [int(degree_channels[2]), 2] + [1] * (len(degree_channels) - 3) diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py new file mode 100644 index 0000000000..766a9d060c --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Degree-wise real Cartesian geometry for DPA4C.""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + TYPE_CHECKING, +) + +import array_api_compat +import numpy as np + +if TYPE_CHECKING: + from collections.abc import ( + Sequence, + ) + + from deepmd.dpmodel.array_api import ( + Array, + ) + +MAX_ANGULAR_DEGREE = 4 +SUPPORTED_CHANNELS = (8, 16, 32, 64, 128) +SUPPORTED_LMAX = (2, 3, 4) + + +def derive_degree_channels(channels: object, lmax: object) -> list[int]: + """Derive the fixed degree-wise channel profile. + + The non-scalar widths follow the scalar width sublinearly: each + non-scalar degree costs ``2 * l + 1`` moment accumulators per channel and + a Gram quadratic in its width, so a linear profile would let the angular + blocks dominate both the reduction payload and the output width. Degree + one takes the geometric mean of the scalar width and the floor, degree two + takes half of that, and degrees three and above keep a single channel. + + Parameters + ---------- + channels + Scalar degree-zero channel width. + lmax + Maximum angular degree. + + Returns + ------- + list[int] + Channel widths for degrees zero through ``lmax``. + + Raises + ------ + TypeError + If either argument is not an integer. + ValueError + If ``channels`` or ``lmax`` is outside the compiled profile set. + """ + if not isinstance(channels, int) or isinstance(channels, bool): + raise TypeError(f"`channels` must be an integer, got {type(channels).__name__}") + if channels not in SUPPORTED_CHANNELS: + raise ValueError( + f"`channels` must be one of {SUPPORTED_CHANNELS}, got {channels}" + ) + if not isinstance(lmax, int) or isinstance(lmax, bool): + raise TypeError(f"`lmax` must be an integer, got {type(lmax).__name__}") + if lmax not in SUPPORTED_LMAX: + raise ValueError(f"`lmax` must be one of {SUPPORTED_LMAX}, got {lmax}") + + # `channels` is a power of two, so the geometric mean is an exact shift. + exponent = channels.bit_length() - 1 + degree_one = max(4, 1 << ((exponent + 1) // 2)) + degree_two = max(4, degree_one >> 1) + return [channels, degree_one, degree_two] + [1] * (lmax - 2) + + +def build_angular_basis(direction: Array, lmax: int) -> Array: + r"""Build normalized real Cartesian harmonics through degree four. + + For vectors :math:`u` and :math:`v`, each degree block obeys + + .. math:: + + B_\ell(u)\cdot B_\ell(v) + =(\lVert u\rVert\lVert v\rVert)^\ell + P_\ell\left( + \frac{u\cdot v}{\lVert u\rVert\lVert v\rVert}\right). + + Parameters + ---------- + direction + Regularized edge directions with shape ``(E, 3)``. + lmax + Maximum angular degree. Supported values are zero through four. + + Returns + ------- + Array + Concatenated degree blocks with shape ``(E, (lmax + 1) ** 2)``. + + Raises + ------ + ValueError + If ``lmax`` is outside the supported range. + """ + if lmax < 0 or lmax > MAX_ANGULAR_DEGREE: + raise ValueError( + f"`lmax` must be between 0 and {MAX_ANGULAR_DEGREE}, got {lmax}" + ) + xp = array_api_compat.array_namespace(direction) + x, y, z = direction[:, 0], direction[:, 1], direction[:, 2] + squared_norm = x * x + y * y + z * z + blocks = [xp.ones_like(x)[:, None]] + if lmax >= 1: + blocks.append(xp.stack([x, y, z], axis=-1)) + if lmax >= 2: + sqrt_three = math.sqrt(3.0) + blocks.append( + xp.stack( + [ + sqrt_three * x * y, + sqrt_three * y * z, + 0.5 * (3.0 * z * z - squared_norm), + sqrt_three * x * z, + 0.5 * sqrt_three * (x * x - y * y), + ], + axis=-1, + ) + ) + if lmax >= 3: + blocks.append( + xp.stack( + [ + math.sqrt(5.0 / 8.0) * y * (3.0 * x * x - y * y), + math.sqrt(15.0) * x * y * z, + math.sqrt(3.0 / 8.0) * y * (5.0 * z * z - squared_norm), + 0.5 * z * (5.0 * z * z - 3.0 * squared_norm), + math.sqrt(3.0 / 8.0) * x * (5.0 * z * z - squared_norm), + 0.5 * math.sqrt(15.0) * z * (x * x - y * y), + math.sqrt(5.0 / 8.0) * x * (x * x - 3.0 * y * y), + ], + axis=-1, + ) + ) + if lmax >= 4: + z_squared = z * z + x2_minus_y2 = x * x - y * y + blocks.append( + xp.stack( + [ + 0.5 * math.sqrt(35.0) * x * y * x2_minus_y2, + 0.25 * math.sqrt(70.0) * y * z * (3.0 * x * x - y * y), + 0.5 * math.sqrt(5.0) * x * y * (7.0 * z_squared - squared_norm), + 0.25 + * math.sqrt(10.0) + * y + * z + * (7.0 * z_squared - 3.0 * squared_norm), + 0.125 + * ( + 35.0 * z_squared * z_squared + - 30.0 * z_squared * squared_norm + + 3.0 * squared_norm * squared_norm + ), + 0.25 + * math.sqrt(10.0) + * x + * z + * (7.0 * z_squared - 3.0 * squared_norm), + 0.25 + * math.sqrt(5.0) + * x2_minus_y2 + * (7.0 * z_squared - squared_norm), + 0.25 * math.sqrt(70.0) * x * z * (x * x - 3.0 * y * y), + 0.125 * math.sqrt(35.0) * (x**4 - 6.0 * x * x * y * y + y**4), + ], + axis=-1, + ) + ) + return xp.concat(blocks, axis=-1) + + +def packed_l2_to_stf(packed: Array) -> Array: + r"""Convert normalized degree-two coefficients to STF matrices. + + The conversion preserves + + .. math:: + + p\cdot q = \operatorname{STF}(p):\operatorname{STF}(q). + + Parameters + ---------- + packed + Degree-two coefficients with shape ``(..., 5)``. + + Returns + ------- + Array + Symmetric-traceless matrices with shape ``(..., 3, 3)``. + """ + xp = array_api_compat.array_namespace(packed) + inv_sqrt_two = 1.0 / math.sqrt(2.0) + inv_sqrt_six = 1.0 / math.sqrt(6.0) + q0, q1, q2, q3, q4 = (packed[..., index] for index in range(5)) + qxy = q0 * inv_sqrt_two + qyz = q1 * inv_sqrt_two + qxz = q3 * inv_sqrt_two + qxx = -q2 * inv_sqrt_six + q4 * inv_sqrt_two + qyy = -q2 * inv_sqrt_six - q4 * inv_sqrt_two + qzz = 2.0 * q2 * inv_sqrt_six + return xp.stack( + [ + xp.stack([qxx, qxy, qxz], axis=-1), + xp.stack([qxy, qyy, qyz], axis=-1), + xp.stack([qxz, qyz, qzz], axis=-1), + ], + axis=-2, + ) + + +def degree_offsets(degree_channels: Sequence[int]) -> tuple[int, ...]: + """Return offsets for the flat degree-wise moment representation. + + Parameters + ---------- + degree_channels + Channel widths for degrees zero through ``lmax``. + + Returns + ------- + tuple[int, ...] + Cumulative offsets with length ``len(degree_channels) + 1``. + """ + offsets = [0] + for degree, width in enumerate(degree_channels): + offsets.append(offsets[-1] + (2 * degree + 1) * int(width)) + return tuple(offsets) + + +def build_moment_indices( + degree_channels: Sequence[int], +) -> tuple[np.ndarray, np.ndarray]: + """Build channel and harmonic indices for one flat edge payload. + + Every degree reads the leading channels of the shared radial map, so the + channel index of a degree is simply ``range(degree_channels[degree])``. + + Parameters + ---------- + degree_channels + Number of channels for degrees zero through ``lmax``. + + Returns + ------- + channel_index + Edge-amplitude channel indices with shape ``(S,)``. + harmonic_index + Packed harmonic indices with shape ``(S,)``, where + ``S = sum((2 * l + 1) * degree_channels[l])``. + """ + channel_index = [] + harmonic_index = [] + for degree, width in enumerate(degree_channels): + for component in range(2 * degree + 1): + channel_index.extend(range(int(width))) + harmonic_index.extend([degree * degree + component] * int(width)) + return ( + np.asarray(channel_index, dtype=np.int64), + np.asarray(harmonic_index, dtype=np.int64), + ) diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py b/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py new file mode 100644 index 0000000000..af521eab84 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Ordered type-pair FiLM cache for DPA4C.""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import array_api_compat + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + NativeOP, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from ..dpa4_nn.mlp import ( + SwiGLUMLP, + resolve_swiglu_hidden_width, +) + + +class OrderedPairFiLM(NativeOP): + r"""Map type embeddings to the ordered finite type-pair cache. + + For center type :math:`a` and neighbor type :math:`b`, the network + evaluates + + .. math:: + + z_{ab}&=[T_a\Vert T_b],\\ + h_{ab}&=\operatorname{SwiGLU}(z_{ab}W_{\rm in}),\\ + [s_{ab},d_{ab},u_{ab}]&=0.1\,h_{ab}W_{\rm out},\\ + \gamma_{ab}&=1+\tanh(s_{ab}),\\ + \beta_{ab}&=T_a+T_b+\tanh(d_{ab}),\\ + U_{ab}&=\tanh(u_{ab}). + + The network is evaluated over the finite type table rather than over graph + edges, so compressed inference stores only the three resulting tables. + Bounding every output keeps the cache well conditioned in ``float32``: + :math:`\gamma` stays in :math:`(0,2)`, and the residual parts of + :math:`\beta` and :math:`U` stay in :math:`(-1,1)`. + + Parameters + ---------- + channels + Type-embedding and FiLM channel width :math:`C`. + radial_modes + Number :math:`R` of shared radial mode profiles each ordered pair + mixes. Zero omits the mixing table. + precision + Parameter precision. + trainable + Whether the pair-encoder weights are trainable. + seed + Random seed. + + Raises + ------ + ValueError + If ``channels`` is not positive or ``radial_modes`` is negative. + """ + + _OUTPUT_SCALE = 0.1 + + def __init__( + self, + channels: int, + radial_modes: int = 0, + *, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + if channels <= 0: + raise ValueError(f"`channels` must be positive, got {channels}") + if radial_modes < 0: + raise ValueError(f"`radial_modes` must be non-negative, got {radial_modes}") + self.channels = int(channels) + self.radial_modes = int(radial_modes) + self.precision = str(precision) + self.trainable = bool(trainable) + input_dim = 2 * self.channels + self.hidden_dim = resolve_swiglu_hidden_width(input_dim) + output_dim = self.channels * (2 + self.radial_modes) + self.network = SwiGLUMLP( + [input_dim, self.hidden_dim, output_dim], + output_scale=self._OUTPUT_SCALE, + precision=self.precision, + trainable=self.trainable, + seed=seed, + ) + + def call(self, type_embedding: Any) -> tuple[Any, Any, Any | None]: + """Build the ordered scale, shift, and mixing tables. + + Parameters + ---------- + type_embedding + Complete type table with shape ``(T + 1, channels)``, where the + trailing row is the zero padding type. + + Returns + ------- + scale + Ordered radial scales with shape ``((T + 1) ** 2, channels)``. + shift + Ordered radial shifts with shape ``((T + 1) ** 2, channels)``. + mixing + Ordered mode-mixing matrices with shape + ``((T + 1) ** 2, channels, radial_modes)``, or ``None`` when + ``radial_modes`` is zero. + """ + xp = array_api_compat.array_namespace(type_embedding) + ntypes = type_embedding.shape[0] + pair_shape = (ntypes, ntypes, self.channels) + pair_input = xp.reshape( + xp.concat( + [ + xp.broadcast_to(type_embedding[:, None, :], pair_shape), + xp.broadcast_to(type_embedding[None, :, :], pair_shape), + ], + axis=-1, + ), + (-1, 2 * self.channels), + ) + logits = self.network.call(pair_input) + + # The output splits into the scale, the shift residual, and the + # flattened mixing matrix, in that order. + shift_end = 2 * self.channels + base_shift = xp.reshape( + type_embedding[:, None, :] + type_embedding[None, :, :], + (-1, self.channels), + ) + return ( + 1.0 + xp.tanh(logits[:, : self.channels]), + base_shift + xp.tanh(logits[:, self.channels : shift_end]), + None + if self.radial_modes == 0 + else xp.reshape( + xp.tanh(logits[:, shift_end:]), + (-1, self.channels, self.radial_modes), + ), + ) + + def serialize(self) -> dict[str, Any]: + """Serialize the ordered pair encoder. + + Returns + ------- + dict[str, Any] + Versioned configuration and pair-encoder parameters. + """ + return { + "@class": "OrderedPairFiLM", + "@version": 1, + "channels": self.channels, + "radial_modes": self.radial_modes, + "precision": self.precision, + "trainable": self.trainable, + "network": self.network.serialize(), + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> OrderedPairFiLM: + """Deserialize an :class:`OrderedPairFiLM`. + + Parameters + ---------- + data + Versioned dictionary produced by :meth:`serialize`. + + Returns + ------- + OrderedPairFiLM + Reconstructed ordered type-pair module. + + Raises + ------ + ValueError + If the payload does not describe an :class:`OrderedPairFiLM`. + """ + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + if data.pop("@class") != "OrderedPairFiLM": + raise ValueError("Invalid serialized class for OrderedPairFiLM") + network = data.pop("network") + obj = cls(**data) + obj.network = SwiGLUMLP.deserialize(network) + return obj diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/readout.py b/deepmd/dpmodel/descriptor/dpa4c_nn/readout.py new file mode 100644 index 0000000000..726f4584d0 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/readout.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Fixed invariant readout for degree-wise DPA4C moments.""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.array_api import ( + xp_asarray_nodetach, +) +from deepmd.dpmodel.utils.network import ( + NativeLayer, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .bispectrum import ( + build_bispectrum_layout, + derive_bispectrum_ranks, +) +from .geometry import ( + degree_offsets, + derive_degree_channels, + packed_l2_to_stf, +) + + +class InvariantReadout(NativeOP): + """Contract degree-wise moments into fixed O(3)-invariant features. + + The readout carries no learned nonlinearity. Degrees one and two first + pass through full-width residual channel maps, after which three fixed + contractions are emitted: the exact channel Gram of every non-scalar + degree, a Cartesian bispectrum over every O(3)-even degree triple, and the + projected quartic ``|Q_b v_a|^2``. The scalar moments are prepended + unchanged. + + Parameters + ---------- + channels + Scalar degree-zero channel width. + lmax + Maximum angular degree. + precision + Parameter precision. + trainable + Whether the channel-alignment and probe projections are trainable. + seed + Random seed reserved for the readout. + """ + + def __init__( + self, + channels: int, + lmax: int, + *, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + self.degree_channels = derive_degree_channels(channels, lmax) + self.bispectrum_ranks = derive_bispectrum_ranks(self.degree_channels) + self.channels = int(channels) + self.lmax = int(lmax) + self.precision = str(precision) + self.trainable = bool(trainable) + self.degree_offsets = degree_offsets(self.degree_channels) + + # === Step 1. Degree-local full channel alignment === + # Only degrees one and two are aligned. Higher degrees carry a single + # channel, for which a residual full-width map is the identity. + alignment_seed = child_seed(seed, 0) + self.channel_alignment = [] + for degree in range(1, 3): + width = self.degree_channels[degree] + self.channel_alignment.append( + NativeLayer( + width, + width, + bias=False, + resnet=True, + precision=self.precision, + seed=child_seed(alignment_seed, degree), + trainable=self.trainable, + ) + ) + + # === Step 2. Exact Gram layout === + # Only the upper triangle is emitted. Scaling the strict off-diagonal + # by sqrt(2) makes the half-vectorization Frobenius isometric. + gram_index_parts = [] + gram_scale_parts = [] + gram_offsets = [0] + for width in self.degree_channels[1:]: + row, column = np.triu_indices(width) + gram_index_parts.append((row * width + column).astype(np.int64)) + gram_scale_parts.append(np.where(row == column, 1.0, math.sqrt(2.0))) + gram_offsets.append(gram_offsets[-1] + row.size) + self.gram_index = np.concatenate(gram_index_parts) + self.gram_scale = np.concatenate(gram_scale_parts).astype( + PRECISION_DICT[self.precision] + ) + self.gram_offsets = tuple(gram_offsets) + + # === Step 3. Bispectrum layout and low-rank probes === + layout = build_bispectrum_layout( + self.lmax, + self.bispectrum_ranks, + ) + self.degree_triples = layout.degree_triples + self.coupling_offsets = layout.coupling_offsets + self.probe_offsets = layout.probe_offsets + self.bispectrum_coupling = layout.coupling.astype( + PRECISION_DICT[self.precision] + ) + self.probe_index = layout.probe_index + self.probe_scale = layout.probe_scale.astype(PRECISION_DICT[self.precision]) + + # A full-rank degree needs no projection. The others start from an + # orthonormal basis with a deterministic sign, so the probes are an + # isometry of a random rank-`K` subspace at initialization. + probe_seed = child_seed(seed, 1) + self.probe_projections = [] + for degree, (width, rank) in enumerate( + zip(self.degree_channels[1:], self.bispectrum_ranks, strict=True), + start=1, + ): + if rank == width: + self.probe_projections.append(None) + continue + degree_seed = child_seed(probe_seed, degree) + projection = NativeLayer( + width, + rank, + bias=False, + precision=self.precision, + seed=degree_seed, + trainable=self.trainable, + ) + # The projection weights are overwritten below, so the basis draw + # reuses the same child seed to stay reproducible. + rng = np.random.default_rng(degree_seed) + orthogonal, triangular = np.linalg.qr( + rng.normal(size=(width, rank)), + mode="reduced", + ) + sign = np.where(np.diag(triangular) < 0.0, -1.0, 1.0) + projection.w = (orthogonal * sign[None, :]).astype( + PRECISION_DICT[self.precision] + ) + self.probe_projections.append(projection) + + self.bispectrum_dim = int(self.probe_index.shape[0]) + self.quartic_dim = self.bispectrum_ranks[0] * self.bispectrum_ranks[1] + + def call(self, moments: Any) -> Any: + """Build invariant features from flat degree-wise moments. + + The output concatenates the scalar moments, the exact aligned Grams in + degree order, the bispectrum in degree-triple order, and the projected + quartic. + + Parameters + ---------- + moments + Flat moment tensor with shape ``(N, S)``, where + ``S = sum((2 * l + 1) * degree_channels[l])``. + + Returns + ------- + Any + Invariant features with shape ``(N, get_dim_out())``. + """ + xp = array_api_compat.array_namespace(moments) + blocks = [ + xp.reshape( + moments[ + :, + self.degree_offsets[degree] : self.degree_offsets[degree + 1], + ], + (moments.shape[0], 2 * degree + 1, self.degree_channels[degree]), + ) + for degree in range(self.lmax + 1) + ] + aligned = list(blocks) + for degree, projection in enumerate(self.channel_alignment, start=1): + aligned[degree] = projection.call(blocks[degree]) + projected = [ + block if projection is None else projection.call(block) + for projection, block in zip( + self.probe_projections, + aligned[1:], + strict=True, + ) + ] + bispectrum, quartic = self.build_bispectrum(projected, xp) + return xp.concat( + [ + blocks[0][:, 0, :], + *self.build_grams(aligned, xp), + *bispectrum, + quartic, + ], + axis=-1, + ) + + def build_grams(self, blocks: list[Any], xp: Any) -> list[Any]: + """Build Frobenius-isometric exact channel Grams. + + Parameters + ---------- + blocks + Aligned degree blocks. Entry ``l`` has shape + ``(N, 2 * l + 1, degree_channels[l])``. + xp + Array namespace associated with ``blocks``. + + Returns + ------- + list[Any] + Upper-triangular Gram blocks for degrees one through ``lmax``. + """ + device = array_api_compat.device(blocks[0]) + gram_index = xp_asarray_nodetach( + xp, + self.gram_index, + device=device, + ) + gram_scale = xp_asarray_nodetach( + xp, + self.gram_scale, + device=device, + ) + parts = [] + for degree, block in enumerate(blocks[1:]): + gram = xp.matmul( + xp.permute_dims(block, (0, 2, 1)), + block, + ) + flat = xp.reshape( + gram, + (block.shape[0], self.degree_channels[degree + 1] ** 2), + ) + start, end = self.gram_offsets[degree : degree + 2] + parts.append( + xp.take(flat, gram_index[start:end], axis=1) + * gram_scale[None, start:end] + ) + return parts + + def build_bispectrum( + self, + projected: list[Any], + xp: Any, + ) -> tuple[list[Any], Any]: + """Contract projected moments with fixed Cartesian Gaunt tensors. + + Parameters + ---------- + projected + Probe-projected degree blocks. Entry ``l - 1`` has shape + ``(N, 2 * l + 1, bispectrum_ranks[l - 1])``. + xp + Array namespace associated with ``projected``. + + Returns + ------- + bispectrum + Independent cubic contractions grouped by degree triple. + quartic + Projected ``|Q_b v_a|^2`` values with shape ``(N, K_2 * K_1)``. + """ + device = array_api_compat.device(projected[0]) + coupling = xp_asarray_nodetach( + xp, + self.bispectrum_coupling, + device=device, + ) + probe_index = xp_asarray_nodetach( + xp, + self.probe_index, + device=device, + ) + probe_scale = xp_asarray_nodetach( + xp, + self.probe_scale, + device=device, + ) + # The 112 triple shares its matrix-vector intermediate with the + # quartic, so it is contracted in closed form rather than through the + # generic Gaunt path. + parts = [] + bispectrum_112, quartic = self.contract_vector_tensor( + projected[0], + projected[1], + xp, + ) + for triple_index, degrees in enumerate(self.degree_triples): + degree_1, degree_2, degree_3 = degrees + if degrees == (1, 1, 2): + full = bispectrum_112 + else: + coupling_start, coupling_end = self.coupling_offsets[ + triple_index : triple_index + 2 + ] + full = self.contract_bispectrum( + xp.reshape( + coupling[coupling_start:coupling_end], + (2 * degree_1 + 1, 2 * degree_2 + 1, 2 * degree_3 + 1), + ), + projected[degree_1 - 1], + projected[degree_2 - 1], + projected[degree_3 - 1], + xp, + ) + probe_start, probe_end = self.probe_offsets[triple_index : triple_index + 2] + parts.append( + xp.take(full, probe_index[probe_start:probe_end], axis=1) + * probe_scale[None, probe_start:probe_end] + ) + return parts, quartic + + def contract_vector_tensor( + self, + vector: Any, + packed_tensor: Any, + xp: Any, + ) -> tuple[Any, Any]: + r"""Contract the ``112`` triple and reuse ``Q_b v_a`` for the quartic. + + Both outputs are built from the same matrix-vector intermediate + :math:`Q_bv_a`, so the quartic costs one extra reduction rather than a + second contraction. + + Parameters + ---------- + vector + Degree-one probes with shape ``(N, 3, K_1)``. + packed_tensor + Degree-two probes with shape ``(N, 5, K_2)``. + xp + Array namespace associated with the probes. + + Returns + ------- + bispectrum_112 + Full ordered cubic contractions with shape ``(N, K_1 * K_1 * K_2)``. + quartic + Values :math:`|Q_bv_a|^2` with shape ``(N, K_2 * K_1)``. + """ + n_nodes = vector.shape[0] + vector_rank = vector.shape[-1] + tensor_rank = packed_tensor.shape[-1] + vectors = xp.permute_dims(vector, (0, 2, 1)) + tensors = packed_l2_to_stf(xp.permute_dims(packed_tensor, (0, 2, 1))) + tensor_vector = xp.reshape( + xp.matmul( + tensors[:, :, None, :, :], + vectors[:, None, :, :, None], + ), + (n_nodes, tensor_rank, vector_rank, 3), + ) + full = xp.matmul( + vectors[:, None, :, :], + xp.permute_dims(tensor_vector, (0, 1, 3, 2)), + ) + full = xp.reshape( + xp.permute_dims(full, (0, 2, 3, 1)), + (n_nodes, vector_rank * vector_rank * tensor_rank), + ) + # The unit-Frobenius 112 Gaunt tensor in this Cartesian convention is + # exactly -v_left^T Q v_right / sqrt(5). + full = full * (-1.0 / math.sqrt(5.0)) + quartic = xp.reshape( + xp.sum(tensor_vector * tensor_vector, axis=-1), + (n_nodes, tensor_rank * vector_rank), + ) + return full, quartic + + def contract_bispectrum( + self, + coupling: Any, + value_1: Any, + value_2: Any, + value_3: Any, + xp: Any, + ) -> Any: + """Contract one angular coupling without backend-specific ``einsum``. + + Parameters + ---------- + coupling + Cartesian Gaunt tensor with shape + ``(2 * l1 + 1, 2 * l2 + 1, 2 * l3 + 1)``. + value_1 + First degree block with shape ``(N, 2 * l1 + 1, K_1)``. + value_2 + Second degree block with shape ``(N, 2 * l2 + 1, K_2)``. + value_3 + Third degree block with shape ``(N, 2 * l3 + 1, K_3)``. + xp + Array namespace associated with the degree blocks. + + Returns + ------- + Any + Ordered contractions with shape ``(N, K_1 * K_2 * K_3)``. + """ + n_nodes = value_1.shape[0] + rank_1 = value_1.shape[-1] + rank_2 = value_2.shape[-1] + rank_3 = value_3.shape[-1] + dim_1, dim_2, dim_3 = coupling.shape + + first = xp.matmul( + xp.permute_dims(value_1, (0, 2, 1)), + xp.reshape(coupling, (dim_1, dim_2 * dim_3)), + ) + first = xp.reshape(first, (n_nodes, rank_1, dim_2, dim_3)) + first = xp.reshape( + xp.permute_dims(first, (0, 1, 3, 2)), + (n_nodes, rank_1 * dim_3, dim_2), + ) + second = xp.matmul(first, value_2) + second = xp.reshape( + second, + (n_nodes, rank_1, dim_3, rank_2), + ) + second = xp.reshape( + xp.permute_dims(second, (0, 1, 3, 2)), + (n_nodes, rank_1 * rank_2, dim_3), + ) + return xp.reshape( + xp.matmul(second, value_3), + (n_nodes, rank_1 * rank_2 * rank_3), + ) + + def get_dim_out(self) -> int: + """Return the geometric output width. + + Returns + ------- + int + Scalar moments, aligned exact Grams, bispectrum probes, and the + projected quartic. + """ + return ( + self.channels + + int(self.gram_index.shape[0]) + + self.bispectrum_dim + + self.quartic_dim + ) + + def serialize(self) -> dict[str, Any]: + """Serialize the invariant readout. + + Returns + ------- + dict[str, Any] + Versioned readout configuration and trainable projections. + """ + return { + "@class": "InvariantReadout", + "@version": 1, + "channels": self.channels, + "lmax": self.lmax, + "precision": self.precision, + "trainable": self.trainable, + "channel_alignment": [ + projection.serialize() for projection in self.channel_alignment + ], + "probe_projections": [ + None if projection is None else projection.serialize() + for projection in self.probe_projections + ], + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> InvariantReadout: + """Deserialize an :class:`InvariantReadout`. + + Parameters + ---------- + data + Versioned dictionary produced by :meth:`serialize`. + + Returns + ------- + InvariantReadout + Reconstructed readout with restored projections. + """ + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + if data.pop("@class") != "InvariantReadout": + raise ValueError("Invalid serialized class for InvariantReadout") + alignment = data.pop("channel_alignment") + probes = data.pop("probe_projections") + obj = cls(**data) + if len(alignment) != len(obj.channel_alignment): + raise ValueError("Serialized alignment projection count is invalid.") + if len(probes) != len(obj.probe_projections): + raise ValueError("Serialized probe projection count is invalid.") + obj.channel_alignment = [ + NativeLayer.deserialize(projection) for projection in alignment + ] + obj.probe_projections = [ + None if projection is None else NativeLayer.deserialize(projection) + for projection in probes + ] + return obj diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 8459d0847f..1abe910a2f 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -498,7 +498,10 @@ def call( ) # [nf, nloc, 3] huber_ncomp = 3 else: - diff_3 = xp.reshape(force_hat - force, (*_node_shape, 3)) + diff_3 = xp.reshape( + force_hat_reshape - force_reshape, + (*_node_shape, 3), + ) norm_2d = xp.reshape( xp.linalg.vector_norm( xp.reshape(diff_3, (-1, 3)), axis=1 @@ -535,7 +538,9 @@ def call( delta=self._huber_delta_force, ) else: - force_diff_3 = xp.reshape(force_hat - force, (-1, 3)) + force_diff_3 = xp.reshape( + force_hat_reshape - force_reshape, (-1, 3) + ) force_diff_norm = xp.reshape( xp.linalg.vector_norm(force_diff_3, axis=1), (-1, 1) ) @@ -554,7 +559,10 @@ def call( if not self.f_use_norm: l1_force_masked = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3) else: - diff_3 = xp.reshape(force_hat - force, (*_node_shape, 3)) + diff_3 = xp.reshape( + force_hat_reshape - force_reshape, + (*_node_shape, 3), + ) norm_2d = xp.reshape( xp.linalg.vector_norm(xp.reshape(diff_3, (-1, 3)), axis=1), _node_shape, @@ -571,7 +579,9 @@ def call( if not self.f_use_norm: l1_force_loss = xp.mean(xp.abs(diff_f)) else: - force_diff_3 = xp.reshape(force_hat - force, (-1, 3)) + force_diff_3 = xp.reshape( + force_hat_reshape - force_reshape, (-1, 3) + ) l1_force_loss = xp.mean( xp.linalg.vector_norm(force_diff_3, axis=1) ) diff --git a/deepmd/kernels/cuda/__init__.py b/deepmd/kernels/cuda/__init__.py index 2ce6f20cc1..6e320d4c30 100644 --- a/deepmd/kernels/cuda/__init__.py +++ b/deepmd/kernels/cuda/__init__.py @@ -16,9 +16,12 @@ / one backward kernel. :mod:`.graph_fitting` Descriptor-agnostic fused energy fitting network on the flat node axis - (cuBLAS GEMMs with fused bias / activation / timestep / residual - epilogues). + (cuBLAS GEMMs with fused bias / activation / residual epilogues). :mod:`.edge_force_virial` Descriptor-agnostic force / atom-virial / per-frame-virial assembly from the per-edge energy gradient. +:mod:`.dpa4c.graph_compress` + DPA4C compressed descriptor: radial spline lookup, two packed moment + reductions, factorized angular feedback, invariant readout, and analytical + edge-vector backward. """ diff --git a/deepmd/kernels/cuda/dpa1/canonical.py b/deepmd/kernels/cuda/dpa1/canonical.py index 856fab953c..113066e2a1 100644 --- a/deepmd/kernels/cuda/dpa1/canonical.py +++ b/deepmd/kernels/cuda/dpa1/canonical.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from deepmd.pt_expt.utils.canonical_graph import ( - DPA1CanonicalGraph, + CanonicalGraph, ) __all__ = [ @@ -203,7 +203,7 @@ def _generic_topology( ) destination_order = torch.arange( source.shape[0], - dtype=source.dtype, + dtype=torch.int64, device=source.device, ) return edge_index, edge_mask, destination_order @@ -277,7 +277,7 @@ def ensure_registered() -> None: def dpa1_canonical_compress_energy_force( desc: Any, fit: Any, - graph: DPA1CanonicalGraph, + graph: CanonicalGraph, atype: torch.Tensor, type_embedding: torch.Tensor, ownership: torch.Tensor, @@ -376,56 +376,38 @@ def dpa1_canonical_compress_energy_force( (int(se.lmax) + 1) ** 2, ) - *hidden, head = fit.nets[0].layers - empty = hidden[0].w.new_empty(0) - weights = [layer.w.contiguous() for layer in hidden] - residuals = [1 if layer.resnet else 0 for layer in hidden] - from deepmd.kernels.triton.dpa1.activation import ( - ACT_CODES, + from deepmd.kernels.cuda.graph_fitting import ( + fitting_operator_arguments, ) + network = fitting_operator_arguments(fit) atom_energy_raw, fitting_saved = torch.ops.deepmd.graph_fitting( descriptor, atype, - weights, - [layer.b.contiguous() if layer.b is not None else empty for layer in hidden], - [ - layer.idt.contiguous() if layer.idt is not None else empty - for layer in hidden - ], - residuals, - head.w.reshape(-1).contiguous(), - ( - head.b.reshape(-1).to(torch.float32).contiguous() - if head.b is not None - else empty - ), + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.head_bias, atom_bias.to(torch.float64).contiguous(), - ACT_CODES[str(hidden[0].activation_function).lower()], + network.activation, ) energy_seed = ownership[:, None].to(atom_energy_raw.dtype) atom_energy = atom_energy_raw * energy_seed - from deepmd.dpmodel.utils.neighbor_graph import ( - frame_id_from_n_node, + from deepmd.kernels.cuda.edge_force_virial import ( + frame_scalar_sum, ) - frame_index = frame_id_from_n_node( - graph.n_node, - n_total=atom_energy.shape[0], - ) - energy = torch.zeros( - graph.n_node.shape[0], - 1, - dtype=atom_energy.dtype, - device=atom_energy.device, - ).index_add_(0, frame_index, atom_energy) + energy = frame_scalar_sum(atom_energy, graph.n_node) del descriptor descriptor_gradient = torch.ops.deepmd.graph_fitting_backward( energy_seed, fitting_saved, - weights, - residuals, - head.w.reshape(-1).contiguous(), + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.activation, ) del fitting_saved edge_gradient = torch.ops.deepmd.dpa1_canonical_compress_backward( diff --git a/deepmd/kernels/cuda/dpa1/graph_compress.py b/deepmd/kernels/cuda/dpa1/graph_compress.py index 27a8c958c8..a05f83b56a 100644 --- a/deepmd/kernels/cuda/dpa1/graph_compress.py +++ b/deepmd/kernels/cuda/dpa1/graph_compress.py @@ -911,8 +911,6 @@ def dpa1_graph_compress_energy_force( lower, upper, table_max, stride0, stride1 = ( float(x) for x in desc.compress_info[0].tolist()[:5] ) - *hidden, head = fit.nets[0].layers - fempty = hidden[0].w.new_empty(0) inverse_stddev = torch.reciprocal(se.stddev[:, 0, :]).contiguous() degree_gain = ( se.adam_degree_gain_raw.to(torch.float32).contiguous() @@ -949,62 +947,40 @@ def dpa1_graph_compress_energy_force( float(se.nnei), (int(se.lmax) + 1) ** 2, ) - weights = [layer.w.contiguous() for layer in hidden] - biases = [ - layer.b.contiguous() if layer.b is not None else fempty for layer in hidden - ] - timesteps = [ - layer.idt.contiguous() if layer.idt is not None else fempty for layer in hidden - ] - residuals = [1 if layer.resnet else 0 for layer in hidden] - head_weight = head.w.reshape(-1).contiguous() - head_bias = ( - head.b.reshape(-1).to(torch.float32).contiguous() - if head.b is not None - else fempty - ) - atom_bias = atom_bias.to(torch.float64).contiguous() - from deepmd.kernels.triton.dpa1.activation import ( - ACT_CODES, + from deepmd.kernels.cuda.graph_fitting import ( + fitting_operator_arguments, ) - activation = ACT_CODES[str(hidden[0].activation_function).lower()] + network = fitting_operator_arguments(fit) + atom_bias = atom_bias.to(torch.float64).contiguous() atom_energy_raw, fitting_saved = torch.ops.deepmd.graph_fitting( descriptor, atype, - weights, - biases, - timesteps, - residuals, - head_weight, - head_bias, + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.head_bias, atom_bias, - activation, + network.activation, ) owned = ownership[:, None].to(atom_energy_raw.dtype) energy_seed = owned atom_energy = atom_energy_raw * owned - from deepmd.dpmodel.utils.neighbor_graph import ( - frame_id_from_n_node, + from deepmd.kernels.cuda.edge_force_virial import ( + frame_scalar_sum, ) - frame_index = frame_id_from_n_node( - graph.n_node, - n_total=atom_energy.shape[0], - ) - energy = torch.zeros( - graph.n_node.shape[0], - 1, - dtype=atom_energy.dtype, - device=atom_energy.device, - ).index_add_(0, frame_index, atom_energy) + energy = frame_scalar_sum(atom_energy, graph.n_node) del descriptor descriptor_gradient = torch.ops.deepmd.graph_fitting_backward( energy_seed, fitting_saved, - weights, - residuals, - head_weight, + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.activation, ) del fitting_saved edge_gradient = torch.ops.deepmd.dpa1_graph_compress_backward( diff --git a/deepmd/kernels/cuda/dpa1/graph_energy_force.py b/deepmd/kernels/cuda/dpa1/graph_energy_force.py index a257a96c92..168c7f3ca8 100644 --- a/deepmd/kernels/cuda/dpa1/graph_energy_force.py +++ b/deepmd/kernels/cuda/dpa1/graph_energy_force.py @@ -94,7 +94,6 @@ def _fake( basis_dim: int, fit_ws: list[torch.Tensor], fit_bs: list[torch.Tensor], - fit_idts: list[torch.Tensor], fit_resnets: list[int], w_head: torch.Tensor, b_head: torch.Tensor, @@ -157,7 +156,6 @@ def _cpu( basis_dim: int, fit_ws: list[torch.Tensor], fit_bs: list[torch.Tensor], - fit_idts: list[torch.Tensor], fit_resnets: list[int], w_head: torch.Tensor, b_head: torch.Tensor, @@ -209,7 +207,6 @@ def _cpu( atype, fit_ws, fit_bs, - fit_idts, fit_resnets, w_head, b_head, @@ -226,7 +223,7 @@ def _cpu( energy = torch.zeros(nf, 1, dtype=atom_e.dtype, device=atom_e.device) energy = energy.index_add(0, frame_id, atom_e) d_grrg = torch.ops.deepmd.graph_fitting_backward( - energy_seed, fit_saved, fit_ws, fit_resnets, w_head + energy_seed, fit_saved, fit_ws, fit_bs, fit_resnets, w_head, fit_act ) del grrg, fit_saved g_e = torch.ops.deepmd.dpa1_graph_descriptor_backward( @@ -371,8 +368,11 @@ def optional(t: torch.Tensor | None) -> torch.Tensor: smooth = 0 w1, w2, w3 = (layer.w.contiguous() for layer in layers) - *hidden, head = fit.nets[0].layers - fempty = hidden[0].w.new_empty(0) + from deepmd.kernels.cuda.graph_fitting import ( + fitting_operator_arguments, + ) + + network = fitting_operator_arguments(fit) if ( graph.destination_order is None or graph.destination_row_ptr is None @@ -423,21 +423,13 @@ def optional(t: torch.Tensor | None) -> torch.Tensor: float(se.env_protection), float(se.nnei), (int(se.lmax) + 1) ** 2, - [layer.w.contiguous() for layer in hidden], - [layer.b.contiguous() if layer.b is not None else fempty for layer in hidden], - [ - layer.idt.contiguous() if layer.idt is not None else fempty - for layer in hidden - ], - [1 if layer.resnet else 0 for layer in hidden], - head.w.reshape(-1).contiguous(), - ( - head.b.reshape(-1).to(torch.float32).contiguous() - if head.b is not None - else fempty - ), + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.head_bias, atom_bias.to(torch.float64).contiguous(), - ACT_CODES[str(hidden[0].activation_function).lower()], + network.activation, node_capacity, do_atomic_virial, ) diff --git a/deepmd/kernels/cuda/dpa4c/__init__.py b/deepmd/kernels/cuda/dpa4c/__init__.py new file mode 100644 index 0000000000..7c70f85de4 --- /dev/null +++ b/deepmd/kernels/cuda/dpa4c/__init__.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Fused compressed CUDA operators for the DPA4C graph lower.""" + +from .canonical import ( + canonical_model_eligible, + dpa4c_canonical_compress_energy_force, +) +from .graph_compress import ( + build_radial_table, + dpa4c_graph_compress, + dpa4c_graph_compress_energy_force, + ef_op_available, + ensure_registered, + mega_eligible, + op_available, +) + +__all__ = [ + "build_radial_table", + "canonical_model_eligible", + "dpa4c_canonical_compress_energy_force", + "dpa4c_graph_compress", + "dpa4c_graph_compress_energy_force", + "ef_op_available", + "ensure_registered", + "mega_eligible", + "op_available", +] diff --git a/deepmd/kernels/cuda/dpa4c/canonical.py b/deepmd/kernels/cuda/dpa4c/canonical.py new file mode 100644 index 0000000000..98cc400724 --- /dev/null +++ b/deepmd/kernels/cuda/dpa4c/canonical.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compact canonical deployment path for compressed DPA4C.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch + +if TYPE_CHECKING: + from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + ) + +_cpu_library: torch.library.Library | None = None + + +def canonical_model_eligible(model: Any) -> bool: + """Return whether a model can use the compact source-only graph ABI.""" + atomic_model = getattr(model, "atomic_model", None) + descriptor = getattr(atomic_model, "descriptor", None) + fitting = getattr(atomic_model, "fitting_net", None) + if descriptor is None or fitting is None: + return False + from deepmd.pt_expt.descriptor.dpa4c import ( + DescrptDPA4C, + ) + + if not isinstance(descriptor, DescrptDPA4C): + return False + if not bool(getattr(descriptor, "compress", False)): + return False + if getattr(descriptor, "exclude_types", None): + return False + if getattr(atomic_model, "pair_excl", None) is not None: + return False + if getattr(atomic_model, "atom_excl", None) is not None: + return False + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + mega_eligible, + ) + from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + ) + + return mega_eligible(descriptor) and fitting_eligible(fitting) + + +def op_available() -> bool: + """Return whether both compact DPA4C descriptor operators are loaded.""" + forward = getattr(torch.ops.deepmd, "dpa4c_canonical_compress", None) + backward = getattr( + torch.ops.deepmd, + "dpa4c_canonical_compress_backward", + None, + ) + backward_inplace = getattr( + torch.ops.deepmd, + "dpa4c_canonical_compress_backward_inplace", + None, + ) + return all( + isinstance(operator, torch._ops.OpOverloadPacket) + for operator in (forward, backward, backward_inplace) + ) + + +def _forward_fake( + edge_vec: torch.Tensor, + source: torch.Tensor, + destination_row_ptr: torch.Tensor, + atype: torch.Tensor, + table: torch.Tensor, + pair_film: torch.Tensor, + pair_mixing: torch.Tensor, + type_embedding: torch.Tensor, + readout_matrices: torch.Tensor, + coupling_meta: torch.Tensor, + coupling_entry: torch.Tensor, + coupling_value: torch.Tensor, + output_mean: torch.Tensor, + output_inv_std: torch.Tensor, + lmax: int, + table_stride: float, + table_max: float, + rcut: float, + eps: float, + degree_floor: float, +) -> tuple[torch.Tensor, torch.Tensor]: + del ( + source, + destination_row_ptr, + table, + pair_film, + pair_mixing, + readout_matrices, + coupling_meta, + coupling_entry, + coupling_value, + output_mean, + output_inv_std, + table_stride, + table_max, + rcut, + eps, + degree_floor, + ) + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + descriptor_profile, + ) + + profile = descriptor_profile(int(type_embedding.shape[1]), int(lmax)) + nodes = atype.shape[0] + descriptor = edge_vec.new_empty(nodes, profile.output_width, dtype=torch.float32) + state = edge_vec.new_empty(nodes, profile.state_width, dtype=torch.float32) + return descriptor, state + + +def _backward_fake( + descriptor_gradient: torch.Tensor, + state: torch.Tensor, + edge_vec: torch.Tensor, + *args: Any, +) -> torch.Tensor: + del descriptor_gradient, state, args + return torch.empty_like(edge_vec) + + +def _backward_inplace_fake( + descriptor_gradient: torch.Tensor, + state: torch.Tensor, + edge_vec: torch.Tensor, + *args: Any, +) -> torch.Tensor: + del descriptor_gradient, state, args + return torch.empty_like(edge_vec) + + +def _energy_gradient_fake( + edge_vec: torch.Tensor, + source: torch.Tensor, + destination_row_ptr: torch.Tensor, + atype: torch.Tensor, + *args: Any, +) -> tuple[torch.Tensor, torch.Tensor]: + del source, destination_row_ptr, args + return ( + edge_vec.new_empty(atype.shape[0], 1, dtype=torch.float64), + torch.empty_like(edge_vec), + ) + + +def _cpu_energy_gradient(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Reference sequence of the fused operator, evaluated in one run.""" + from deepmd.kernels.cuda.graph_fitting import _cpu_backward as fitting_backward + from deepmd.kernels.cuda.graph_fitting import _cpu_forward as fitting_forward + + descriptor_args = args[:_DESCRIPTOR_ARGUMENT_COUNT] + ws, bs, resnets, w_head, b_head, bias_atom_e, act, seed, _tile = args[ + _DESCRIPTOR_ARGUMENT_COUNT: + ] + descriptor, state = _cpu_forward(*descriptor_args) + atype = descriptor_args[3] + energy, saved = fitting_forward( + descriptor, atype, ws, bs, resnets, w_head, b_head, bias_atom_e, act + ) + gradient = fitting_backward( + seed.reshape(-1, 1), saved, ws, bs, resnets, w_head, act + ) + edge_gradient = _cpu_backward(gradient, state, *descriptor_args) + return energy, edge_gradient + + +def _generic_topology( + source: torch.Tensor, + destination_row_ptr: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Materialize the generic topology for the small CPU trace sample.""" + physical_edge_count = int(destination_row_ptr[-1].item()) + node_count = destination_row_ptr.shape[0] - 1 + destination = torch.repeat_interleave( + torch.arange(node_count, dtype=torch.int64, device=source.device), + destination_row_ptr[1:] - destination_row_ptr[:-1], + output_size=physical_edge_count, + ) + source_i64 = source.to(torch.int64) + destination_storage = torch.zeros_like(source_i64) + destination_storage[:physical_edge_count] = destination + edge_index = torch.stack((source_i64, destination_storage)) + edge_mask = ( + torch.arange(source.shape[0], dtype=torch.int64, device=source.device) + < physical_edge_count + ) + destination_order = torch.arange( + source.shape[0], + dtype=torch.int64, + device=source.device, + ) + return edge_index, edge_mask, destination_order + + +# The compact ABI drops the three topology tensors of the generic ABI and its +# leading ``canonical`` flag; the remaining trailing scalars are identical. +_CANONICAL_SCALAR_COUNT = 6 + +#: Leading arguments of the fused operator that describe the descriptor: +#: ``edge_vec`` plus the compact topology, the compression artifacts and the +#: six trailing geometry scalars. +_DESCRIPTOR_ARGUMENT_COUNT = 20 + + +def _cpu_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: + from deepmd.kernels.cuda.dpa4c.graph_compress import _cpu_forward as generic_forward + + edge_vec, source, destination_row_ptr, atype, *tail = args + edge_index, edge_mask, destination_order = _generic_topology( + source, + destination_row_ptr, + ) + return generic_forward( + edge_vec, + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + atype, + *tail[:-_CANONICAL_SCALAR_COUNT], + True, + *tail[-_CANONICAL_SCALAR_COUNT:], + ) + + +def _cpu_backward(*args: Any) -> torch.Tensor: + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + _cpu_backward as generic_backward, + ) + + descriptor_gradient, state, edge_vec, source, destination_row_ptr, atype, *tail = ( + args + ) + edge_index, edge_mask, destination_order = _generic_topology( + source, + destination_row_ptr, + ) + return generic_backward( + descriptor_gradient, + state, + edge_vec, + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + atype, + *tail[:-_CANONICAL_SCALAR_COUNT], + True, + *tail[-_CANONICAL_SCALAR_COUNT:], + ) + + +def _cpu_backward_inplace(*args: Any) -> torch.Tensor: + edge_gradient = _cpu_backward(*args) + args[1].zero_() + return edge_gradient + + +def ensure_registered() -> None: + """Register fake and CPU implementations for compact DPA4C operators.""" + global _cpu_library + if _cpu_library is not None or not op_available(): + return + torch.library.register_fake("deepmd::dpa4c_canonical_compress")(_forward_fake) + torch.library.register_fake("deepmd::dpa4c_canonical_compress_backward")( + _backward_fake + ) + torch.library.register_fake("deepmd::dpa4c_canonical_compress_energy_gradient")( + _energy_gradient_fake + ) + torch.library.register_fake("deepmd::dpa4c_canonical_compress_backward_inplace")( + _backward_inplace_fake + ) + _cpu_library = torch.library.Library("deepmd", "IMPL") + _cpu_library.impl("dpa4c_canonical_compress", _cpu_forward, "CPU") + _cpu_library.impl( + "dpa4c_canonical_compress_backward", + _cpu_backward, + "CPU", + ) + _cpu_library.impl( + "dpa4c_canonical_compress_energy_gradient", + _cpu_energy_gradient, + "CPU", + ) + _cpu_library.impl( + "dpa4c_canonical_compress_backward_inplace", + _cpu_backward_inplace, + "CPU", + ) + + +def dpa4c_canonical_compress_energy_force( + descriptor: Any, + fitting: Any, + graph: CanonicalGraph, + atype: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate compressed DPA4C from a compact canonical edge stream. + + The compact ABI carries only source indices and CSR row pointers, so the + descriptor operator addresses edges by position and omits every mask, + padding-type, and table-tail check that the cutoff-compacted graph already + guarantees. The backward reuses the forward state in place. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor. + fitting + Eligible pt_expt energy fitting network. + graph + Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` and + whose ``destination_row_ptr`` has shape ``(N + 1,)`` in ``int64``. + atype + Flat node atom types with shape ``(N,)``. + ownership + Boolean mask selecting energy-contributing nodes with shape ``(N,)``. + atom_bias + Combined atomic energy bias with shape ``(ntypes,)`` in eV. + do_atomic_virial + Whether to return per-node virials. + + Returns + ------- + energy + Per-frame energy with shape ``(F, 1)`` in eV, fp64. + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, fp64. + force + Per-node force with shape ``(N, 3)`` in eV/Å, fp32. + virial + Per-frame virial with shape ``(F, 3, 3)`` in eV, fp32. + atom_virial + Per-node virial with shape ``(N, 3, 3)`` in eV, or an empty tensor. + + Raises + ------ + ValueError + If the model or the compiled operators do not support the compact path. + """ + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + compressed_operator_arguments, + mega_eligible, + ) + from deepmd.kernels.cuda.edge_force_virial import ( + canonical_edge_force_virial, + canonical_op_available, + ) + from deepmd.kernels.cuda.edge_force_virial import ( + ensure_registered as ensure_force_registered, + ) + from deepmd.kernels.cuda.edge_force_virial import ( + frame_scalar_sum, + ) + from deepmd.kernels.cuda.graph_fitting import ( + ensure_registered as ensure_fitting_registered, + ) + from deepmd.kernels.cuda.graph_fitting import ( + fitting_operator_arguments, + node_tile, + ) + + ensure_registered() + ensure_fitting_registered() + ensure_force_registered() + if not mega_eligible(descriptor) or not canonical_op_available(): + raise ValueError("model is not eligible for compact canonical DPA4C inference") + + network = fitting_operator_arguments(fitting) + atom_energy_raw, edge_gradient = ( + torch.ops.deepmd.dpa4c_canonical_compress_energy_gradient( + graph.edge_vec, + graph.source, + graph.destination_row_ptr, + atype, + *compressed_operator_arguments(descriptor), + int(descriptor.lmax), + *descriptor._compression_scalars, + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.head_bias, + atom_bias.to(torch.float64).contiguous(), + network.activation, + ownership.to(torch.float64).reshape(-1).contiguous(), + node_tile(), + ) + ) + atom_energy = atom_energy_raw * ownership[:, None].to(atom_energy_raw.dtype) + energy = frame_scalar_sum(atom_energy, graph.n_node) + force, atom_virial, virial = canonical_edge_force_virial( + edge_gradient, + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + graph.n_node, + atype.shape[0], + do_atomic_virial, + ) + return energy, atom_energy, force, virial, atom_virial diff --git a/deepmd/kernels/cuda/dpa4c/graph_compress.py b/deepmd/kernels/cuda/dpa4c/graph_compress.py new file mode 100644 index 0000000000..f35f25f13f --- /dev/null +++ b/deepmd/kernels/cuda/dpa4c/graph_compress.py @@ -0,0 +1,1395 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Bindings and immutable artifacts for compressed degree-wise DPA4C. + +The CUDA operator evaluates the complete DPA4C graph descriptor: the tabulated +radial branch and its shared mode profiles, the ordered PairFiLM amplitude, one +destination reduction that produces both envelope masses and every degree-wise +moment, the invariant readout, and the fixed output calibration. Its analytical +backward runs one node-readout VJP and one edge recomputation scan. + +The radial table stores the learned scalar-distance maps + +.. math:: + + r\mapsto g_c(r)=\operatorname{RadialMLP}(\operatorname{RBF}(r))_c,\qquad + r\mapsto q_\rho(r)=\operatorname{ModeHead}(\operatorname{RadialMLP} + _{\rm hidden}(\operatorname{RBF}(r)))_\rho, + +so its width is ``channels + radial_modes``. The fixed C³ envelope remains +analytical and gates the complete FiLM amplitude exactly once; the non-scalar +moments carry a second explicit envelope factor. Ordered scale, shift, and +mode-mixing caches, the readout projections, and the sparse angular coupling +tables are materialized once when compression is enabled. + +Degrees one and two dominate the readout and are contracted in closed form. +Degrees three and four carry a single channel each, so their couplings are +driven by a compact sparse Cartesian Gaunt table rather than by specialized +code. +""" + +from __future__ import ( + annotations, +) + +import copy +import math +from dataclasses import ( + dataclass, +) +from typing import ( + Any, +) + +import numpy as np +import torch + +from deepmd.dpmodel.descriptor.dpa4c_nn import ( + build_angular_basis, + build_bispectrum_layout, + derive_bispectrum_ranks, + derive_degree_channels, + packed_l2_to_stf, +) + +__all__ = [ + "build_compression_artifacts", + "build_radial_table", + "coupling_records", + "descriptor_profile", + "dpa4c_graph_compress", + "dpa4c_graph_compress_energy_force", + "ef_op_available", + "ensure_registered", + "fitting_energy_and_gradient", + "mega_eligible", + "op_available", +] + +SUPPORTED_CHANNELS = (8, 16, 32, 64, 128) +SUPPORTED_LMAX = (2, 3, 4) +SUPPORTED_RADIAL_MODES = (0, 2, 4, 8) + +# Degrees one and two carry the wide channel blocks and are contracted by +# specialized closed forms in every backend; the remaining triples run through +# the shared sparse coupling path. +_CLOSED_FORM_TRIPLES = ((1, 1, 2), (2, 2, 2)) + +# Lebedev quadrature leaves rounding noise around eighteen orders of magnitude +# below the retained coupling entries, far below the fp32 working precision. +_COUPLING_TOLERANCE = 1.0e-9 + +# Unit-Frobenius Cartesian normalizations of the two closed-form couplings. +_INV_SQRT_FIVE = 1.0 / math.sqrt(5.0) +_BIS222_SCALE = math.sqrt(12.0 / 35.0) + + +def op_available() -> bool: + """Return whether the compiled DPA4C compressed operator is loaded.""" + op = getattr(torch.ops.deepmd, "dpa4c_graph_compress", None) + return isinstance(op, torch._ops.OpOverloadPacket) + + +def ef_op_available() -> bool: + """Return whether the descriptor, fitting, and force operators are loaded.""" + return ( + op_available() + and isinstance( + getattr(torch.ops.deepmd, "graph_fitting", None), + torch._ops.OpOverloadPacket, + ) + and isinstance( + getattr(torch.ops.deepmd, "edge_force_virial", None), + torch._ops.OpOverloadPacket, + ) + ) + + +def mega_eligible(descriptor: Any) -> bool: + """Return whether the descriptor has a compiled fp32 specialization. + + Parameters + ---------- + descriptor + Descriptor instance, which need not be a DPA4C. + + Returns + ------- + bool + Whether the compiled operator covers this configuration. + """ + return ( + int(getattr(descriptor, "channels", -1)) in SUPPORTED_CHANNELS + and int(getattr(descriptor, "lmax", -1)) in SUPPORTED_LMAX + and int(getattr(descriptor, "radial_modes", -1)) in SUPPORTED_RADIAL_MODES + and str(getattr(descriptor, "precision", "")).lower() in ("float32", "single") + ) + + +# === Compiled profile === + + +@dataclass(frozen=True) +class DescriptorProfile: + r"""Describe every width the compiled operator derives from its inputs. + + Attributes + ---------- + channels + Scalar degree-zero width :math:`C_0`. + lmax + Maximum angular degree. + degree_channels + Channel width of each degree from zero through ``lmax``. + ranks + Bispectrum probe rank of each degree from one through ``lmax``. + moment_width + Flat moment width :math:`S=\\sum_\\ell(2\\ell+1)C_\\ell`. + output_width + Invariant descriptor width consumed by the fitting network. + gram_base + Descriptor coordinate of the first degree-one Gram entry. + bispectrum_base + Descriptor coordinate of the first bispectrum entry. + """ + + channels: int + lmax: int + degree_channels: tuple[int, ...] + ranks: tuple[int, ...] + moment_width: int + output_width: int + gram_base: int + bispectrum_base: int + + @property + def state_width(self) -> int: + """Return the saved-state width: the moments plus both normalizers.""" + return self.moment_width + 2 + + +def descriptor_profile(channels: int, lmax: int) -> DescriptorProfile: + """Derive every compiled width from the two structural parameters. + + Parameters + ---------- + channels + Scalar degree-zero width. + lmax + Maximum angular degree. + + Returns + ------- + DescriptorProfile + Widths and descriptor-block offsets shared by all backends. + """ + degree_channels = tuple(derive_degree_channels(int(channels), int(lmax))) + ranks = tuple(derive_bispectrum_ranks(degree_channels)) + moment_width = sum( + (2 * degree + 1) * width for degree, width in enumerate(degree_channels) + ) + gram_total = sum(width * (width + 1) // 2 for width in degree_channels[1:]) + layout = build_bispectrum_layout(int(lmax), ranks) + bispectrum_dim = int(layout.probe_index.shape[0]) + gram_base = degree_channels[0] + bispectrum_base = gram_base + gram_total + # Geometric block, the two moment divisors, then the center type tail. + output_width = ( + bispectrum_base + bispectrum_dim + ranks[0] * ranks[1] + 2 + degree_channels[0] + ) + return DescriptorProfile( + channels=int(channels), + lmax=int(lmax), + degree_channels=degree_channels, + ranks=ranks, + moment_width=moment_width, + output_width=output_width, + gram_base=gram_base, + bispectrum_base=bispectrum_base, + ) + + +@dataclass(frozen=True) +class CouplingRecord: + r"""Describe one sparse angular coupling consumed by the operator. + + Attributes + ---------- + degrees + Degree triple :math:`(\\ell_1,\\ell_2,\\ell_3)`. + nonzero_begin + First entry index of the coupling nonzeros. + nonzero_count + Number of coupling nonzeros. + probe_begin + First entry index of the packed probe coordinates. + probe_count + Number of emitted probe contractions. + coordinate + Descriptor coordinate of the first emitted contraction. + """ + + degrees: tuple[int, int, int] + nonzero_begin: int + nonzero_count: int + probe_begin: int + probe_count: int + coordinate: int + + +def coupling_records( + channels: int, + lmax: int, +) -> tuple[list[CouplingRecord], np.ndarray, np.ndarray]: + """Build the sparse coupling tables for degrees beyond the closed forms. + + Each record addresses one contiguous run of coupling nonzeros followed by + one contiguous run of probe coordinates inside the two shared flat arrays. + A coupling nonzero packs its harmonic components as + ``m1 | m2 << 8 | m3 << 16`` and carries the Gaunt value; a probe coordinate + packs its channel indices the same way and carries the isometric scale. + + Parameters + ---------- + channels + Scalar degree-zero width. + lmax + Maximum angular degree. + + Returns + ------- + records + One record per non-closed-form degree triple. + entry + Packed component and channel coordinates with shape ``(M,)``, int32. + value + Coupling values and probe scales with shape ``(M,)``, float32. + """ + profile = descriptor_profile(channels, lmax) + layout = build_bispectrum_layout(profile.lmax, profile.ranks) + records: list[CouplingRecord] = [] + entries: list[int] = [] + values: list[float] = [] + for index, degrees in enumerate(layout.degree_triples): + if degrees in _CLOSED_FORM_TRIPLES: + continue + degree_1, degree_2, degree_3 = degrees + start, end = layout.coupling_offsets[index : index + 2] + coupling = layout.coupling[start:end].reshape( + 2 * degree_1 + 1, + 2 * degree_2 + 1, + 2 * degree_3 + 1, + ) + nonzero_begin = len(entries) + for component in np.argwhere(np.abs(coupling) > _COUPLING_TOLERANCE): + first, second, third = (int(value) for value in component) + entries.append(first | (second << 8) | (third << 16)) + values.append(float(coupling[first, second, third])) + nonzero_count = len(entries) - nonzero_begin + + rank_2 = profile.ranks[degree_2 - 1] + rank_3 = profile.ranks[degree_3 - 1] + start, end = layout.probe_offsets[index : index + 2] + probe_begin = len(entries) + for probe, scale in zip( + layout.probe_index[start:end], + layout.probe_scale[start:end], + strict=True, + ): + third = int(probe) % rank_3 + second = (int(probe) // rank_3) % rank_2 + first = int(probe) // (rank_3 * rank_2) + entries.append(first | (second << 8) | (third << 16)) + values.append(float(scale)) + records.append( + CouplingRecord( + degrees=degrees, + nonzero_begin=nonzero_begin, + nonzero_count=nonzero_count, + probe_begin=probe_begin, + probe_count=end - start, + coordinate=profile.bispectrum_base + int(start), + ) + ) + return ( + records, + np.asarray(entries, dtype=np.int32), + np.asarray(values, dtype=np.float32), + ) + + +def _coupling_meta(records: list[CouplingRecord]) -> np.ndarray: + """Flatten coupling records into the int32 metadata table.""" + return np.asarray( + [ + [ + record.degrees[0], + record.degrees[1], + record.degrees[2], + record.nonzero_begin, + record.nonzero_count, + record.probe_begin, + record.probe_count, + record.coordinate, + ] + for record in records + ], + dtype=np.int32, + ).reshape(len(records), 8) + + +# === Immutable artifacts === + + +def _quintic_coefficients( + values: torch.Tensor, + first: torch.Tensor, + second: torch.Tensor, + stride: float, +) -> torch.Tensor: + """Build C²-matching quintic Hermite coefficients. + + Parameters + ---------- + values + Function values with shape ``(S + 1, W)``. + first + First derivatives with shape ``(S + 1, W)``. + second + Second derivatives with shape ``(S + 1, W)``. + stride + Uniform interval width in Å. + The interval row is split into a leading quartet block holding + ``[c0, c1, c2, c3]`` of every channel and a trailing pair block holding + ``[c4, c5]``. Both blocks are naturally aligned for one 128-bit and one + 64-bit load, which evaluates the spline in two memory instructions instead + of the three a plain channel-major layout would need, at identical traffic. + + Returns + ------- + torch.Tensor + Coefficients with shape ``(S, 6 * W)``: ``4 * W`` quartet entries + followed by ``2 * W`` pair entries. + """ + left_value, right_value = values[:-1], values[1:] + left_first, right_first = first[:-1], first[1:] + left_second, right_second = second[:-1], second[1:] + delta = right_value - left_value + h = float(stride) + c0 = left_value + c1 = left_first + c2 = 0.5 * left_second + c3 = ( + 20.0 * delta + - (8.0 * right_first + 12.0 * left_first) * h + - (3.0 * left_second - right_second) * h * h + ) / (2.0 * h**3) + c4 = ( + -30.0 * delta + + (14.0 * right_first + 16.0 * left_first) * h + + (3.0 * left_second - 2.0 * right_second) * h * h + ) / (2.0 * h**4) + c5 = ( + 12.0 * delta + - 6.0 * (right_first + left_first) * h + + (right_second - left_second) * h * h + ) / (2.0 * h**5) + intervals = values.shape[0] - 1 + return torch.cat( + [ + torch.stack([c0, c1, c2, c3], dim=-1).reshape(intervals, -1), + torch.stack([c4, c5], dim=-1).reshape(intervals, -1), + ], + dim=-1, + ) + + +def build_radial_table( + descriptor: Any, + stride: float = 0.002, +) -> tuple[torch.Tensor, torch.Tensor]: + """Tabulate the composed DPA4C radial branch and its mode profiles. + + The table width is ``channels + radial_modes``: the leading block holds the + shared radial map read by every degree, and the trailing block holds the + shared mode profiles that each ordered type pair mixes. + + Parameters + ---------- + descriptor + pt_expt DPA4C descriptor whose DPA4 radial modules define the table. + stride + Uniform distance spacing in Å. + + Returns + ------- + table + Quintic coefficients with shape ``(S, 6 * (channels + radial_modes))``, + fp32. + info + CPU metadata ``[stride, table_max, rcut, eps, degree_floor]``, fp64. + + Raises + ------ + ValueError + If ``stride`` is not positive. + """ + if stride <= 0.0: + raise ValueError(f"`stride` must be positive, got {stride}") + sample_parameter = next(descriptor.radial_embedding.parameters()) + device = sample_parameter.device + radial_basis = copy.deepcopy(descriptor.radial_basis).to( + device=device, + dtype=torch.float64, + ) + radial_embedding = copy.deepcopy(descriptor.radial_embedding).to( + device=device, + dtype=torch.float64, + ) + radial_mode_head = ( + None + if descriptor.radial_mode_head is None + else copy.deepcopy(descriptor.radial_mode_head).to( + device=device, + dtype=torch.float64, + ) + ) + interval_count = math.ceil(float(descriptor.rcut) / float(stride)) + table_max = interval_count * float(stride) + distance = torch.arange( + interval_count + 1, + dtype=torch.float64, + device=device, + ) + distance = (distance * float(stride)).requires_grad_(True) + + # === Step 1. Evaluate every tabulated distance map at the table knots === + hidden = radial_embedding.call_hidden(radial_basis(distance[:, None])) + values = radial_embedding.call_output(hidden) + if radial_mode_head is not None: + values = torch.cat([values, radial_mode_head.call(hidden)], dim=-1) + + # === Step 2. Differentiate each output channel independently === + first_columns = [] + second_columns = [] + for channel in range(int(values.shape[1])): + (first_channel,) = torch.autograd.grad( + values[:, channel].sum(), + distance, + create_graph=True, + retain_graph=True, + ) + (second_channel,) = torch.autograd.grad( + first_channel.sum(), + distance, + retain_graph=True, + ) + first_columns.append(first_channel) + second_columns.append(second_channel) + first = torch.stack(first_columns, dim=-1) + second = torch.stack(second_columns, dim=-1) + + # === Step 3. Convert knot data to the runtime spline layout === + table = _quintic_coefficients(values, first, second, float(stride)) + table = table.detach().to(torch.float32).contiguous() + info = torch.tensor( + [ + float(stride), + table_max, + float(descriptor.rcut), + float(descriptor._EPS), + float(descriptor._DEGREE_NORM_FLOOR), + ], + dtype=torch.float64, + device="cpu", + ) + return table, info + + +def build_compression_artifacts( + descriptor: Any, + stride: float = 0.002, +) -> dict[str, torch.Tensor]: + """Build all immutable tensors consumed by compressed inference. + + Parameters + ---------- + descriptor + Evaluated pt_expt DPA4C descriptor. + stride + Uniform radial spline spacing in Å. + + Returns + ------- + dict[str, torch.Tensor] + Radial spline coefficients, metadata, ordered FiLM, mixing and type + caches, padded readout projections, sparse angular couplings, and the + output calibration. + + Raises + ------ + ValueError + If the descriptor is not an fp32 model with a compiled specialization. + """ + sample_parameter = next(descriptor.parameters()) + if sample_parameter.dtype != torch.float32: + raise ValueError( + "DPA4C compressed CUDA requires descriptor precision `float32`, " + f"got {sample_parameter.dtype}" + ) + if not mega_eligible(descriptor): + raise ValueError( + "DPA4C compressed CUDA supports channels " + f"{SUPPORTED_CHANNELS} with lmax {SUPPORTED_LMAX} and " + f"radial_modes {SUPPORTED_RADIAL_MODES}, got " + f"channels={descriptor.channels}, lmax={descriptor.lmax}, " + f"radial_modes={descriptor.radial_modes}" + ) + device = sample_parameter.device + profile = descriptor_profile(descriptor.channels, descriptor.lmax) + table, info = build_radial_table(descriptor, stride) + + with torch.no_grad(): + type_embedding = descriptor.type_embedding.call().to( + device=device, + dtype=torch.float32, + ) + pair_scale, pair_shift, pair_mixing = descriptor.pair_film.call(type_embedding) + pair_film = torch.stack((pair_scale, pair_shift), dim=-1) + # The mode axis is innermost so that the coefficients a lane needs for + # one channel arrive in one or two vector loads. + mixing = ( + torch.zeros(0, dtype=torch.float32, device=device) + if pair_mixing is None + else pair_mixing.to(torch.float32) + ) + readout_matrices = _build_readout_matrices(descriptor, profile, device) + output_mean = descriptor.mean.to(device=device, dtype=torch.float32) + output_inv_std = torch.reciprocal( + descriptor.stddev.to(device=device, dtype=torch.float32) + ) + + records, coupling_entry, coupling_value = coupling_records( + profile.channels, + profile.lmax, + ) + to_device = {"device": device} + return { + "data": table, + "info": info, + "pair_film": pair_film.detach().contiguous(), + "pair_mixing": mixing.detach().contiguous(), + "type_embedding": type_embedding.detach().contiguous(), + "readout_matrices": readout_matrices, + "coupling_meta": torch.as_tensor( + _coupling_meta(records), + dtype=torch.int32, + **to_device, + ), + "coupling_entry": torch.as_tensor( + coupling_entry, + dtype=torch.int32, + **to_device, + ), + "coupling_value": torch.as_tensor( + coupling_value, + dtype=torch.float32, + **to_device, + ), + "output_mean": output_mean.detach().contiguous(), + "output_inv_std": output_inv_std.detach().contiguous(), + } + + +def _build_readout_matrices( + descriptor: Any, + profile: DescriptorProfile, + device: torch.device, +) -> torch.Tensor: + """Pack the degree-one and degree-two readout projections. + + The kernel reads the residual channel alignment and the probe projection of + each wide degree together with its transpose, so all eight matrices are + stored in one square block padded to the degree-one width. Degrees three + and above carry a single channel, for which both maps are the identity. + + Parameters + ---------- + descriptor + Evaluated pt_expt DPA4C descriptor. + profile + Compiled profile of this descriptor. + device + Device that receives the packed table. + + Returns + ------- + torch.Tensor + Packed projections with shape ``(8, C_1, C_1)``, fp32. + """ + width = profile.degree_channels[1] + packed = torch.zeros(8, width, width, dtype=torch.float32, device=device) + + def residual(layer: Any, size: int) -> torch.Tensor: + return layer.w.to(torch.float32) + torch.eye( + size, + dtype=torch.float32, + device=device, + ) + + def probe(layer: Any, size: int) -> torch.Tensor: + if layer is None: + return torch.eye(size, dtype=torch.float32, device=device) + return layer.w.to(torch.float32) + + alignment_one = residual( + descriptor.readout.channel_alignment[0], + profile.degree_channels[1], + ) + alignment_two = residual( + descriptor.readout.channel_alignment[1], + profile.degree_channels[2], + ) + probe_one = probe( + descriptor.readout.probe_projections[0], + profile.degree_channels[1], + ) + probe_two = probe( + descriptor.readout.probe_projections[1], + profile.degree_channels[2], + ) + for index, matrix in enumerate( + ( + alignment_one, + alignment_one.T, + alignment_two, + alignment_two.T, + probe_one, + probe_one.T, + probe_two, + probe_two.T, + ) + ): + packed[index, : matrix.shape[0], : matrix.shape[1]] = matrix + return packed.detach().contiguous() + + +# === Reference implementation === + + +def _table_lookup( + table: torch.Tensor, + radius: torch.Tensor, + stride: float, + table_max: float, + width: int, +) -> torch.Tensor: + """Evaluate a uniform quintic table with a clamped high-distance tail.""" + coordinate = radius.clamp(min=0.0, max=table_max) + index = torch.floor(coordinate / stride).to(torch.int64) + index = index.clamp(max=table.shape[0] - 1) + dx = (coordinate - index.to(coordinate.dtype) * stride)[:, None] + row = table[index] + quartet = row[:, : 4 * width].reshape(-1, width, 4) + pair = row[:, 4 * width :].reshape(-1, width, 2) + return ( + quartet[..., 0] + + ( + quartet[..., 1] + + ( + quartet[..., 2] + + (quartet[..., 3] + (pair[..., 0] + pair[..., 1] * dx) * dx) * dx + ) + * dx + ) + * dx + ) + + +def _c3_envelope(radius: torch.Tensor, rcut: float) -> torch.Tensor: + """Evaluate the fixed exponent-five DPA4 C³ envelope.""" + u = ((float(rcut) - radius) / float(rcut)).clamp(0.0, 1.0) + x = 1.0 - u + series = 1.0 + x * (4.0 + x * (10.0 + x * (20.0 + 35.0 * x))) + return u**4 * series + + +def _half_gram(value: torch.Tensor) -> torch.Tensor: + """Return the Frobenius-isometric upper-triangular channel Gram.""" + width = value.shape[-1] + row, column = torch.triu_indices(width, width, device=value.device) + scale = torch.where( + row == column, + torch.ones((), dtype=value.dtype, device=value.device), + torch.full((), math.sqrt(2.0), dtype=value.dtype, device=value.device), + ) + return (value.transpose(1, 2) @ value)[:, row, column] * scale + + +def _contract_coupling( + coupling: torch.Tensor, + first: torch.Tensor, + second: torch.Tensor, + third: torch.Tensor, +) -> torch.Tensor: + """Contract one angular coupling into ordered probe combinations. + + Parameters + ---------- + coupling + Gaunt tensor with shape ``(2 * l1 + 1, 2 * l2 + 1, 2 * l3 + 1)``. + first + Degree block with shape ``(N, 2 * l1 + 1, K_1)``. + second + Degree block with shape ``(N, 2 * l2 + 1, K_2)``. + third + Degree block with shape ``(N, 2 * l3 + 1, K_3)``. + + Returns + ------- + torch.Tensor + Ordered contractions with shape ``(N, K_1 * K_2 * K_3)``. + """ + nodes = first.shape[0] + rank_1, rank_2, rank_3 = first.shape[-1], second.shape[-1], third.shape[-1] + dim_1, dim_2, dim_3 = coupling.shape + value = first.transpose(1, 2) @ coupling.reshape(dim_1, dim_2 * dim_3) + value = value.reshape(nodes, rank_1, dim_2, dim_3).transpose(2, 3) + value = (value.reshape(nodes, rank_1 * dim_3, dim_2) @ second).reshape( + nodes, + rank_1, + dim_3, + rank_2, + ) + value = value.transpose(2, 3).reshape(nodes, rank_1 * rank_2, dim_3) + return (value @ third).reshape(nodes, rank_1 * rank_2 * rank_3) + + +def _cpu_descriptor( + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + atype: torch.Tensor, + table: torch.Tensor, + pair_film: torch.Tensor, + pair_mixing: torch.Tensor, + type_embedding: torch.Tensor, + readout_matrices: torch.Tensor, + coupling_meta: torch.Tensor, + coupling_entry: torch.Tensor, + coupling_value: torch.Tensor, + output_mean: torch.Tensor, + output_inv_std: torch.Tensor, + canonical: bool, + lmax: int, + table_stride: float, + table_max: float, + rcut: float, + eps: float, + degree_floor: float, +) -> torch.Tensor: + """Reference implementation of the compressed DPA4C descriptor.""" + del destination_order, destination_row_ptr, canonical, coupling_meta + compute = edge_vec.to(torch.float32) + source, destination = edge_index[0].to(torch.long), edge_index[1].to(torch.long) + node_count = atype.shape[0] + channels = type_embedding.shape[1] + type_count = type_embedding.shape[0] + profile = descriptor_profile(int(channels), int(lmax)) + radial_modes = 0 if pair_mixing.numel() == 0 else int(pair_mixing.shape[2]) + + # === Step 1. Build the masked edge geometry === + radius = torch.sqrt((compute * compute).sum(dim=-1) + float(eps) ** 2) + direction = compute / radius[:, None] + center_type = atype[destination] + neighbor_type = atype[source] + mask = edge_mask & (neighbor_type < type_count - 1) & (center_type < type_count - 1) + maskf = mask.to(compute.dtype) + envelope = _c3_envelope(radius, float(rcut)) * maskf + + # === Step 2. Evaluate the ordered FiLM amplitude === + tabulated = _table_lookup( + table.to(compute.device), + radius, + float(table_stride), + float(table_max), + channels + radial_modes, + ) + pair_index = center_type * type_count + neighbor_type + film = pair_film[pair_index] + amplitude = tabulated[:, :channels] * film[..., 0] + film[..., 1] + if radial_modes > 0: + amplitude = amplitude + torch.einsum( + "ecr,er->ec", + pair_mixing[pair_index], + tabulated[:, channels:], + ) + amplitude = amplitude * envelope[:, None] + basis = build_angular_basis(direction, profile.lmax) * maskf[:, None] + + # === Step 3. Reduce both envelope masses and every degree-wise moment === + envelope_squared = envelope * envelope + scalar_mass = torch.zeros( + node_count, + dtype=compute.dtype, + device=compute.device, + ).index_add_(0, destination, envelope_squared) + angular_mass = torch.zeros( + node_count, + dtype=compute.dtype, + device=compute.device, + ).index_add_(0, destination, envelope_squared * envelope_squared) + scalar_divisor = torch.sqrt(scalar_mass + float(degree_floor))[:, None] + angular_divisor = torch.sqrt(angular_mass + float(degree_floor))[:, None] + scalar_norm = torch.reciprocal(scalar_divisor) + angular_norm = torch.reciprocal(angular_divisor)[:, :, None] + + blocks = [ + torch.zeros( + node_count, + 1, + channels, + dtype=compute.dtype, + device=compute.device, + ).index_add_(0, destination, amplitude[:, None, :]) + * scalar_norm[:, None] + ] + for degree, width in enumerate(profile.degree_channels[1:], start=1): + payload = ( + basis[:, degree**2 : (degree + 1) ** 2, None] + * amplitude[:, None, :width] + * envelope[:, None, None] + ) + blocks.append( + torch.zeros( + node_count, + 2 * degree + 1, + width, + dtype=compute.dtype, + device=compute.device, + ).index_add_(0, destination, payload) + * angular_norm + ) + + # === Step 4. Align, project, and contract the invariant readout === + width_one, width_two = profile.degree_channels[1], profile.degree_channels[2] + rank_one, rank_two = profile.ranks[0], profile.ranks[1] + aligned = list(blocks) + aligned[1] = blocks[1] @ readout_matrices[0, :width_one, :width_one] + aligned[2] = blocks[2] @ readout_matrices[2, :width_two, :width_two] + probes = list(aligned[1:]) + probes[0] = aligned[1] @ readout_matrices[4, :width_one, :rank_one] + probes[1] = aligned[2] @ readout_matrices[6, :width_two, :rank_two] + + vectors = probes[0].transpose(1, 2) + tensors = packed_l2_to_stf(probes[1].transpose(1, 2)) + tensor_vector = (tensors[:, :, None] @ vectors[:, None, :, :, None]).squeeze(-1) + parts: dict[int, torch.Tensor] = {} + for record in coupling_records(int(channels), profile.lmax)[0]: + degree_1, degree_2, degree_3 = record.degrees + start = record.nonzero_begin + components = coupling_entry[start : start + record.nonzero_count] + coupling = torch.zeros( + 2 * degree_1 + 1, + 2 * degree_2 + 1, + 2 * degree_3 + 1, + dtype=compute.dtype, + device=compute.device, + ) + coupling[ + components & 0xFF, + (components >> 8) & 0xFF, + (components >> 16) & 0xFF, + ] = coupling_value[start : start + record.nonzero_count].to(compute.dtype) + full = _contract_coupling( + coupling, + probes[degree_1 - 1], + probes[degree_2 - 1], + probes[degree_3 - 1], + ) + start = record.probe_begin + selection = coupling_entry[start : start + record.probe_count] + rank_2 = profile.ranks[degree_2 - 1] + rank_3 = profile.ranks[degree_3 - 1] + flat = ( + ((selection & 0xFF) * rank_2 + ((selection >> 8) & 0xFF)) * rank_3 + + ((selection >> 16) & 0xFF) + ).to(torch.long) + parts[record.coordinate] = full[:, flat] * coupling_value[ + start : start + record.probe_count + ].to(compute.dtype) + + parts[profile.bispectrum_base] = _closed_form_112( + vectors, + tensor_vector, + rank_one, + rank_two, + ) + parts[_closed_form_222_coordinate(profile)] = _closed_form_222(tensors) + + # === Step 5. Assemble and calibrate the invariant output === + quartic = (tensor_vector * tensor_vector).sum(dim=-1).flatten(start_dim=1) + descriptor = torch.cat( + [ + blocks[0][:, 0, :], + *[_half_gram(block) for block in aligned[1:]], + *[parts[key] for key in sorted(parts)], + quartic, + scalar_divisor, + angular_divisor, + type_embedding[atype], + ], + dim=-1, + ) + return (descriptor - output_mean[None, :]) * output_inv_std[None, :] + + +def _closed_form_112( + vectors: torch.Tensor, + tensor_vector: torch.Tensor, + rank_one: int, + rank_two: int, +) -> torch.Tensor: + """Contract the ``112`` triple from the shared ``Q_b v_a`` intermediate.""" + values = [] + for first in range(rank_one): + for second in range(first, rank_one): + scale = 1.0 if first == second else math.sqrt(2.0) + for index in range(rank_two): + values.append( + -scale + * _INV_SQRT_FIVE + * (vectors[:, first] * tensor_vector[:, index, second]).sum(dim=-1) + ) + return torch.stack(values, dim=-1) + + +def _closed_form_222(tensors: torch.Tensor) -> torch.Tensor: + """Contract the symmetric ``222`` triple over the two degree-two probes. + + One operator ordering suffices because the factors are symmetric, so + ``tr(ABC) = tr((ABC)^T) = tr(CBA) = tr(ACB)``. + """ + values = [] + for first, second, third, scale in ( + (0, 0, 0, 1.0), + (0, 0, 1, math.sqrt(3.0)), + (0, 1, 1, math.sqrt(3.0)), + (1, 1, 1, 1.0), + ): + product = tensors[:, first] @ tensors[:, second] @ tensors[:, third] + values.append( + -_BIS222_SCALE + * scale + * torch.linalg.diagonal(product, dim1=-2, dim2=-1).sum(-1) + ) + return torch.stack(values, dim=-1) + + +def _closed_form_222_coordinate(profile: DescriptorProfile) -> int: + """Return the descriptor coordinate of the ``222`` bispectrum block.""" + layout = build_bispectrum_layout(profile.lmax, profile.ranks) + index = layout.degree_triples.index((2, 2, 2)) + return profile.bispectrum_base + int(layout.probe_offsets[index]) + + +# === Custom-operator registration === + + +def _cpu_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: + """CPU custom-op implementation returning descriptor and opaque state.""" + descriptor = _cpu_descriptor(*args) + profile = descriptor_profile(int(args[9].shape[1]), int(args[17])) + state = torch.zeros( + descriptor.shape[0], + profile.state_width, + dtype=descriptor.dtype, + device=descriptor.device, + ) + return descriptor, state + + +def _forward_fake( + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + atype: torch.Tensor, + table: torch.Tensor, + pair_film: torch.Tensor, + pair_mixing: torch.Tensor, + type_embedding: torch.Tensor, + readout_matrices: torch.Tensor, + coupling_meta: torch.Tensor, + coupling_entry: torch.Tensor, + coupling_value: torch.Tensor, + output_mean: torch.Tensor, + output_inv_std: torch.Tensor, + canonical: bool, + lmax: int, + table_stride: float, + table_max: float, + rcut: float, + eps: float, + degree_floor: float, +) -> tuple[torch.Tensor, torch.Tensor]: + del ( + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + table, + pair_film, + pair_mixing, + readout_matrices, + coupling_meta, + coupling_entry, + coupling_value, + output_mean, + output_inv_std, + canonical, + table_stride, + table_max, + rcut, + eps, + degree_floor, + ) + profile = descriptor_profile(int(type_embedding.shape[1]), int(lmax)) + descriptor = torch.empty( + atype.shape[0], + profile.output_width, + dtype=torch.float32, + device=edge_vec.device, + ) + state = torch.empty( + atype.shape[0], + profile.state_width, + dtype=torch.float32, + device=edge_vec.device, + ) + return descriptor, state + + +def _backward_fake( + descriptor_gradient: torch.Tensor, + state: torch.Tensor, + edge_vec: torch.Tensor, + *args: Any, +) -> torch.Tensor: + del descriptor_gradient, state, args + return torch.empty_like(edge_vec) + + +def _cpu_backward( + descriptor_gradient: torch.Tensor, + state: torch.Tensor, + edge_vec: torch.Tensor, + *args: Any, +) -> torch.Tensor: + del state + if edge_vec.shape[0] == 0: + return torch.zeros_like(edge_vec) + value = edge_vec.detach().clone().requires_grad_(True) + with torch.enable_grad(): + descriptor = _cpu_descriptor(value, *args) + (gradient,) = torch.autograd.grad( + (descriptor * descriptor_gradient.to(descriptor.dtype)).sum(), + value, + ) + return gradient.to(edge_vec.dtype) + + +def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: + ctx.save_for_backward(output[1], *inputs[:16]) + ctx.scalars = inputs[16:] + ctx.mark_non_differentiable(output[1]) + ctx.set_materialize_grads(False) + + +def _backward( + ctx: Any, + descriptor_gradient: torch.Tensor, + state_gradient: torch.Tensor | None, +) -> tuple: + del state_gradient + tensors = ctx.saved_tensors + edge_gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( + descriptor_gradient, + tensors[0], + *tensors[1:], + *ctx.scalars, + ) + return (edge_gradient,) + (None,) * 22 + + +_cpu_library: torch.library.Library | None = None + + +def ensure_registered() -> None: + """Register fake, CPU, and autograd implementations once.""" + global _cpu_library + if _cpu_library is not None or not op_available(): + return + torch.library.register_fake("deepmd::dpa4c_graph_compress")(_forward_fake) + torch.library.register_fake("deepmd::dpa4c_graph_compress_backward")(_backward_fake) + torch.library.register_autograd( + "deepmd::dpa4c_graph_compress", + _backward, + setup_context=_setup_context, + ) + _cpu_library = torch.library.Library("deepmd", "IMPL") + _cpu_library.impl("dpa4c_graph_compress", _cpu_forward, "CPU") + _cpu_library.impl( + "dpa4c_graph_compress_backward", + _cpu_backward, + "CPU", + ) + + +def compressed_operator_arguments(descriptor: Any) -> tuple: + """Return the immutable operator arguments of a compressed descriptor. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor. + + Returns + ------- + tuple + Radial table, ordered caches, readout projections, coupling tables, + output calibration, and the trailing scalar configuration. + """ + return ( + descriptor.compress_data, + descriptor.compress_pair_film, + descriptor.compress_pair_mixing, + descriptor.compress_type_embedding, + descriptor.compress_readout_matrices, + descriptor.compress_coupling_meta, + descriptor.compress_coupling_entry, + descriptor.compress_coupling_value, + descriptor.compress_output_mean, + descriptor.compress_output_inv_std, + ) + + +def dpa4c_graph_compress( + descriptor: Any, + graph: Any, + atype: torch.Tensor, +) -> torch.Tensor: + """Evaluate the compressed DPA4C graph descriptor. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor. + graph + NeighborGraph with destination CSR topology. + atype + Flat node types with shape ``(N,)``. + + Returns + ------- + torch.Tensor + Degree-wise invariant descriptor with shape + ``(N, descriptor.get_dim_out())``, fp32. + + Raises + ------ + ValueError + If the graph lacks destination CSR topology. + """ + ensure_registered() + if graph.destination_order is None or graph.destination_row_ptr is None: + raise ValueError("DPA4C compressed CUDA requires destination CSR topology") + descriptor_output, _state = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec.contiguous(), + graph.edge_index.contiguous(), + graph.edge_mask.contiguous(), + graph.destination_order.contiguous(), + graph.destination_row_ptr.contiguous(), + atype.contiguous(), + *compressed_operator_arguments(descriptor), + bool(graph.destination_sorted), + int(descriptor.lmax), + *descriptor._compression_scalars, + ) + return descriptor_output.to(graph.edge_vec.dtype) + + +def dpa4c_graph_compress_energy_force( + descriptor: Any, + fitting: Any, + graph: Any, + atype: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + node_capacity: int, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate compressed DPA4C energy, force, and virial without a tape. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor. + fitting + Eligible pt_expt energy fitting network. + graph + NeighborGraph with destination and source CSR topology. + atype + Flat node atom types with shape ``(N,)``. + ownership + Boolean mask selecting energy-contributing nodes with shape ``(N,)``. + atom_bias + Combined atomic energy bias with shape ``(ntypes,)``. + node_capacity + Force-scatter node capacity. + do_atomic_virial + Whether to return per-node virials. + + Returns + ------- + energy + Per-frame energy with shape ``(F, 1)``, fp64. + atom_energy + Per-node energy with shape ``(N, 1)``, fp64. + force + Per-node force with shape ``(N, 3)``, fp32. + virial + Per-frame virial with shape ``(F, 3, 3)``, fp32. + atom_virial + Per-node virial with shape ``(N, 3, 3)`` or an empty tensor. + + Raises + ------ + ValueError + If the graph lacks destination or source CSR topology. + """ + from deepmd.kernels.cuda.edge_force_virial import ( + edge_force_virial, + ) + from deepmd.kernels.cuda.edge_force_virial import ( + ensure_registered as ensure_force_registered, + ) + from deepmd.kernels.cuda.graph_fitting import ( + ensure_registered as ensure_fitting_registered, + ) + + ensure_registered() + ensure_fitting_registered() + ensure_force_registered() + if ( + graph.destination_order is None + or graph.destination_row_ptr is None + or graph.source_order is None + or graph.source_row_ptr is None + ): + raise ValueError( + "DPA4C compressed energy-force inference requires destination " + "and source CSR topology" + ) + + operator_args = ( + graph.edge_index.contiguous(), + graph.edge_mask.contiguous(), + graph.destination_order.contiguous(), + graph.destination_row_ptr.contiguous(), + atype.contiguous(), + *compressed_operator_arguments(descriptor), + bool(graph.destination_sorted), + int(descriptor.lmax), + *descriptor._compression_scalars, + ) + edge_vec = graph.edge_vec.to(torch.float32).contiguous() + node_descriptor, state = torch.ops.deepmd.dpa4c_graph_compress( + edge_vec, + *operator_args, + ) + + energy, atom_energy, descriptor_gradient = fitting_energy_and_gradient( + fitting, + node_descriptor, + atype, + ownership, + atom_bias, + graph.n_node, + ) + del node_descriptor + edge_gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( + descriptor_gradient, + state, + edge_vec, + *operator_args, + ) + force, atom_virial, virial = edge_force_virial( + edge_gradient, + edge_vec, + graph.edge_index, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + graph.n_node, + node_capacity, + do_atomic_virial, + ) + return energy, atom_energy, force, virial, atom_virial + + +def fitting_energy_and_gradient( + fitting: Any, + node_descriptor: torch.Tensor, + atype: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + n_node_per_frame: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the fused fitting network and its descriptor cotangent. + + Parameters + ---------- + fitting + Eligible pt_expt energy fitting network. + node_descriptor + Invariant descriptor with shape ``(N, D)``, fp32. + atype + Flat node atom types with shape ``(N,)``. + ownership + Boolean mask selecting energy-contributing nodes with shape ``(N,)``. + atom_bias + Combined atomic energy bias with shape ``(ntypes,)``. + n_node_per_frame + Node count of each frame with shape ``(F,)``. + + Returns + ------- + energy + Per-frame energy with shape ``(F, 1)``, fp64. + atom_energy + Per-node energy with shape ``(N, 1)``, fp64. + descriptor_gradient + Cotangent of the invariant descriptor with shape ``(N, D)``. + """ + from deepmd.kernels.cuda.edge_force_virial import ( + frame_scalar_sum, + ) + from deepmd.kernels.cuda.graph_fitting import ( + energy_and_input_gradient, + ) + + atom_energy_raw, descriptor_gradient = energy_and_input_gradient( + fitting, + node_descriptor, + atype, + ownership, + atom_bias, + ) + atom_energy = atom_energy_raw * ownership[:, None].to(atom_energy_raw.dtype) + energy = frame_scalar_sum(atom_energy, n_node_per_frame) + return energy, atom_energy, descriptor_gradient diff --git a/deepmd/kernels/cuda/edge_force_virial.py b/deepmd/kernels/cuda/edge_force_virial.py index 22cbe58007..e7ea7b772d 100644 --- a/deepmd/kernels/cuda/edge_force_virial.py +++ b/deepmd/kernels/cuda/edge_force_virial.py @@ -34,6 +34,8 @@ "canonical_op_available", "edge_force_virial", "ensure_registered", + "frame_scalar_sum", + "frame_scalar_sum_available", "op_available", ] @@ -50,6 +52,65 @@ def canonical_op_available() -> bool: return isinstance(op, torch._ops.OpOverloadPacket) +def frame_scalar_sum_available() -> bool: + """Whether the C++ ``deepmd::frame_scalar_sum`` op is loaded.""" + op = getattr(torch.ops.deepmd, "frame_scalar_sum", None) + return isinstance(op, torch._ops.OpOverloadPacket) + + +def _frame_scalar_sum_fake( + node_scalar: torch.Tensor, + n_node_per_frame: torch.Tensor, +) -> torch.Tensor: + return node_scalar.new_empty(n_node_per_frame.shape[0], 1) + + +def _frame_scalar_sum_cpu( + node_scalar: torch.Tensor, + n_node_per_frame: torch.Tensor, +) -> torch.Tensor: + offsets = torch.cat( + [ + torch.zeros(1, dtype=torch.int64, device=n_node_per_frame.device), + torch.cumsum(n_node_per_frame.to(torch.int64), 0), + ] + ) + return torch.stack( + [ + node_scalar[offsets[frame] : offsets[frame + 1]].sum(0) + for frame in range(n_node_per_frame.shape[0]) + ] + ) + + +def frame_scalar_sum( + node_scalar: torch.Tensor, + n_node_per_frame: torch.Tensor, +) -> torch.Tensor: + """Sum a node-major scalar over the node segment of each frame. + + Parameters + ---------- + node_scalar : torch.Tensor + Per-node scalar with shape ``(N, 1)``. + n_node_per_frame : torch.Tensor + Node count of each frame with shape ``(F,)``. The frames occupy + contiguous spans of the node axis in this order. + + Returns + ------- + torch.Tensor + Per-frame total with shape ``(F, 1)`` and the input dtype. + + Notes + ----- + Nodes past ``sum(n_node_per_frame)`` are padding of the flat node axis and + contribute to no frame. + """ + ensure_registered() + return torch.ops.deepmd.frame_scalar_sum(node_scalar, n_node_per_frame) + + def _fake( g_e: torch.Tensor, edge_vec: torch.Tensor, @@ -164,7 +225,7 @@ def _canonical_cpu( edge_mask, torch.arange( edge_vec.shape[0], - dtype=source_order.dtype, + dtype=torch.int64, device=edge_vec.device, ), destination_row_ptr, @@ -192,8 +253,12 @@ def ensure_registered() -> None: torch.library.register_fake("deepmd::canonical_edge_force_virial")( _canonical_fake ) + if frame_scalar_sum_available(): + torch.library.register_fake("deepmd::frame_scalar_sum")(_frame_scalar_sum_fake) _cpu_library = torch.library.Library("deepmd", "IMPL") _cpu_library.impl("edge_force_virial", _cpu, "CPU") + if frame_scalar_sum_available(): + _cpu_library.impl("frame_scalar_sum", _frame_scalar_sum_cpu, "CPU") if canonical_op_available(): _cpu_library.impl( "canonical_edge_force_virial", diff --git a/deepmd/kernels/cuda/graph_fitting.py b/deepmd/kernels/cuda/graph_fitting.py index fb8b7624c5..2243dba2e8 100644 --- a/deepmd/kernels/cuda/graph_fitting.py +++ b/deepmd/kernels/cuda/graph_fitting.py @@ -4,10 +4,10 @@ The CUDA operator ``deepmd::graph_fitting`` (see ``source/op/pt/graph_fitting.cu``) evaluates the whole energy fitting network on the flat node axis -- cuBLAS GEMMs with the bias / activation / -timestep / residual epilogues fused into single elementwise kernels -- and -returns the per-atom energy in fp64. The registered backward chains the layer -dgrads from the saved activation derivatives, exposing the descriptor -gradient that the force / virial assembly differentiates through. +residual epilogues fused into single elementwise kernels -- and returns the +per-atom energy in fp64. The registered backward chains the layer dgrads from +the saved pre-activations, exposing the descriptor gradient that the force / +virial assembly differentiates through. The operator is descriptor-agnostic: any graph-lowered energy model whose fitting is a plain MLP over the flat node axis (see @@ -16,26 +16,33 @@ Usage and pitfalls ------------------ -* The forward's second output packs the layer activation derivatives as one - flat buffer of ``adot chunks`` (the activations themselves are a forward-only - transient); it is an autograd save, never a user-facing value, and receives - no gradient (``set_materialize_grads(False)``). -* The backward infers the node count and descriptor width from the saved - derivative buffer and first weight. It deliberately does not retain the - descriptor tensor, allowing inference memory planners to reuse descriptor - storage for its gradient after the fitting forward. +* The forward's second output packs the layer pre-activations as one flat + buffer, chunk per layer, written by the GEMMs themselves; the activations are + a forward-only transient in a two-slot ping-pong. The buffer is an autograd + save, never a user-facing value, and is marked non-differentiable. +* The backward re-derives each activation derivative from the saved + pre-activation and the layer bias, so it takes the biases and the activation + code as arguments. It infers the node count and descriptor width from the + saved buffer and first weight, deliberately not retaining the descriptor + tensor, which allows inference memory planners to reuse descriptor storage + for its gradient after the fitting forward. * The head bias is passed as a device tensor, not a Python float: reading a value host-side (``.item()``) inside the dispatch path would fail under symbolic tracing (``GuardOnDataDependentSymNode``) and force a GPU sync per step in eager mode. -* ``Tensor[]`` op inputs (weights / biases / timesteps) are pytree list - nodes: the backward must return a matching ``list`` of ``None`` for each, - while ``int[]`` inputs are single leaves taking a single ``None``. -* The eligibility gate (:func:`fitting_eligible`) requires float4-aligned - hidden widths, one tanh / silu activation on every hidden layer, a linear - scalar head and no frame / atomic parameters. +* ``Tensor[]`` op inputs (weights / biases) are pytree list nodes: the + backward must return a matching ``list`` of ``None`` for each, while + ``int[]`` inputs are single leaves taking a single ``None``. +* The operator represents exactly the networks :func:`fitting_eligible` + accepts and silently computes something else for any other network, so + every entry point builds its arguments through + :func:`fitting_operator_arguments`, which validates before it converts. """ +import os +from dataclasses import ( + dataclass, +) from typing import ( Any, ) @@ -47,12 +54,29 @@ ) __all__ = [ + "FittingArguments", + "energy_and_input_gradient", "ensure_registered", "fitting_eligible", + "fitting_operator_arguments", "graph_fitting", + "node_tile", "op_available", ] +#: Nodes evaluated per run of a tiled inference pipeline. The tile bounds +#: every node-scale allocation that a run retires, and at this size the extra +#: launches are a fraction of a percent of a saturated step while the GEMM +#: shape stays far above the width at which cuBLAS loses efficiency. +#: ``DP_NODE_TILE`` overrides it; zero evaluates the whole node axis at once. +_DEFAULT_TILE = 131072 + + +def node_tile() -> int: + """Return the configured node tile of the tiled inference pipeline.""" + value = os.environ.get("DP_NODE_TILE") + return _DEFAULT_TILE if value is None else int(value) + def op_available() -> bool: """Whether the C++ ``deepmd::graph_fitting`` op is loaded.""" @@ -64,10 +88,18 @@ def fitting_eligible(fit: Any) -> bool: """Whether the fused fitting operator can serve this network. Requires a single mixed-types energy net with tanh / silu on every hidden - layer, fp32 weights and hidden-layer parameters, a linear scalar head - without timestep or residual, float4-aligned hidden widths, and no frame / - atomic parameters, case embedding or type exclusion. Hidden identity - residuals are supported; width-doubling residuals use the reference path. + layer, fp32 weights and biases, no layer timestep, a linear scalar head + without residual, float4-aligned hidden widths of at most 4096, and no + frame / atomic parameters, case embedding or type exclusion. Hidden + identity residuals are supported; width-doubling residuals use the + reference path. + + A per-layer timestep is a supported configuration of the reference + network but not of this operator. Nothing prevents it technically -- the + scale is a length-``dout`` vector, so the backward could form + ``act'(pre + b) * idt`` from the same saved pre-activation -- but the + deployed models do not use one, and carrying it would add a load and a + multiply per element to the forward and to both backward epilogues. Parameters ---------- @@ -96,10 +128,12 @@ def fitting_eligible(fit: Any) -> bool: return False if head.w.shape[1] != 1 or head.idt is not None or head.resnet: return False + if any(layer.idt is not None for layer in hidden): + return False tensors = [ tensor for layer in layers - for tensor in (layer.w, layer.b, layer.idt) + for tensor in (layer.w, layer.b) if tensor is not None ] if any(tensor.dtype != torch.float32 for tensor in tensors): @@ -108,7 +142,88 @@ def fitting_eligible(fit: Any) -> bool: layer.resnet and layer.w.shape[1] == 2 * layer.w.shape[0] for layer in hidden ): return False - return all(int(layer.w.shape[1]) % 4 == 0 for layer in hidden) + # The elementwise epilogues map one float4 lane of a row to threadIdx.x, + # so a row must be non-empty and fit the 1024-thread block limit at four + # values per lane. + return all( + int(layer.w.shape[1]) % 4 == 0 and 0 < int(layer.w.shape[1]) <= 4096 + for layer in hidden + ) + + +@dataclass(frozen=True) +class FittingArguments: + """Fused operator arguments derived from one fitting network. + + Attributes + ---------- + weights : list[torch.Tensor] + Hidden layer weights, each with shape (din, dout). + biases : list[torch.Tensor] + Hidden layer biases with shape (dout,), empty where a layer has none. + residuals : list[int] + One flag per hidden layer marking an identity residual. + head_weight : torch.Tensor + Flattened linear head weight with shape (width,). + head_bias : torch.Tensor + Scalar head bias with shape (1,), empty where the head has none. + activation : int + Hidden activation code, see ``ACT_CODES``. + """ + + weights: list[torch.Tensor] + biases: list[torch.Tensor] + residuals: list[int] + head_weight: torch.Tensor + head_bias: torch.Tensor + activation: int + + +def fitting_operator_arguments(fit: Any) -> FittingArguments: + """Convert a fitting network into fused operator arguments. + + The operator represents exactly the networks :func:`fitting_eligible` + accepts. For anything else it does not fail; it evaluates a different + network, dropping whatever it cannot represent. Validation therefore + belongs at every boundary that converts a module into operator arguments, + which is what this function is. + + Parameters + ---------- + fit : EnergyFittingNet + The pt_expt fitting module. + + Returns + ------- + FittingArguments + Contiguous tensors and codes ready for ``deepmd::graph_fitting``. + + Raises + ------ + ValueError + If the fused operator cannot reproduce this network. + """ + if not fitting_eligible(fit): + raise ValueError( + "the fused fitting operator cannot reproduce this network; test " + "fitting_eligible() first and fall back to the reference path" + ) + *hidden, head = fit.nets[0].layers + empty = hidden[0].w.new_empty(0) + return FittingArguments( + weights=[layer.w.contiguous() for layer in hidden], + biases=[ + layer.b.contiguous() if layer.b is not None else empty for layer in hidden + ], + residuals=[1 if layer.resnet else 0 for layer in hidden], + head_weight=head.w.reshape(-1).contiguous(), + head_bias=( + head.b.reshape(-1).to(torch.float32).contiguous() + if head.b is not None + else empty + ), + activation=ACT_CODES[str(hidden[0].activation_function).lower()], + ) # ====================================================================== @@ -119,7 +234,6 @@ def _forward_fake( atype: torch.Tensor, ws: list[torch.Tensor], bs: list[torch.Tensor], - idts: list[torch.Tensor], resnets: list[int], w_head: torch.Tensor, b_head: torch.Tensor, @@ -134,12 +248,30 @@ def _forward_fake( ) +def _energy_gradient_fake( + x: torch.Tensor, + atype: torch.Tensor, + ws: list[torch.Tensor], + bs: list[torch.Tensor], + resnets: list[int], + w_head: torch.Tensor, + b_head: torch.Tensor, + bias_atom_e: torch.Tensor, + act: int, + seed: torch.Tensor, + tile: int, +) -> torch.Tensor: + return x.new_empty(x.shape[0], 1, dtype=torch.float64) + + def _backward_fake( d_e: torch.Tensor, saved: torch.Tensor, ws: list[torch.Tensor], + bs: list[torch.Tensor], resnets: list[int], w_head: torch.Tensor, + act: int, ) -> torch.Tensor: total_width = sum(int(w.shape[1]) for w in ws) n_node = saved.shape[0] // total_width @@ -147,62 +279,66 @@ def _backward_fake( def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: - x, atype, ws, bs, idts, resnets, w_head, b_head, bias_atom_e, act = inputs + x, atype, ws, bs, resnets, w_head, b_head, bias_atom_e, act = inputs _e, saved = output - ctx.save_for_backward(saved, w_head, *ws) + # The saved buffer is an internal autograd artifact; the backward consumes + # it as data and produces no cotangent for it. + ctx.mark_non_differentiable(saved) + ctx.save_for_backward(saved, w_head, *ws, *bs) ctx.n_layers = len(ws) ctx.resnets = resnets + ctx.act = act ctx.set_materialize_grads(False) def _backward(ctx: Any, d_e: torch.Tensor, d_saved: Any) -> tuple: - saved, w_head, *ws = ctx.saved_tensors + saved, w_head, *rest = ctx.saved_tensors + ws, bs = rest[: ctx.n_layers], rest[ctx.n_layers :] d_x = torch.ops.deepmd.graph_fitting_backward( - d_e, saved, list(ws), ctx.resnets, w_head + d_e, saved, list(ws), list(bs), ctx.resnets, w_head, ctx.act ) none_list = [None] * ctx.n_layers - return (d_x, None, none_list, none_list, none_list, None, None, None, None, None) + return (d_x, None, none_list, none_list, None, None, None, None, None) # ====================================================================== # CPU reference implementations # ====================================================================== +def _activation_derivative(pre: torch.Tensor, act: int) -> torch.Tensor: + """Return the activation derivative at the given pre-activation.""" + if act == 0: + return 1.0 - torch.tanh(pre) ** 2 + sigmoid = torch.sigmoid(pre) + return sigmoid * (1.0 + pre * (1.0 - sigmoid)) + + def _cpu_forward( x: torch.Tensor, atype: torch.Tensor, ws: list[torch.Tensor], bs: list[torch.Tensor], - idts: list[torch.Tensor], resnets: list[int], w_head: torch.Tensor, b_head: torch.Tensor, bias_atom_e: torch.Tensor, act: int, ) -> tuple[torch.Tensor, torch.Tensor]: - adots = [] + pres = [] cur = x.to(torch.float32) - for w, b, idt, res in zip(ws, bs, idts, resnets): + for w, b, res in zip(ws, bs, resnets, strict=True): pre = cur @ w + pres.append(pre) if b.numel(): pre = pre + b a = torch.tanh(pre) if act == 0 else torch.nn.functional.silu(pre) - if act == 0: - adot = 1.0 - a * a - else: - s = torch.sigmoid(pre) - adot = s * (1.0 + pre * (1.0 - s)) - if idt.numel(): - a = a * idt - adot = adot * idt - adots.append(adot) cur = a + cur if (res and w.shape[0] == w.shape[1]) else a e = (cur @ w_head[:, None]).to(torch.float64) if b_head.numel(): e = e + b_head.to(torch.float64) e = e + bias_atom_e[atype][:, None] - # Chunk layout mirrors the CUDA op: adot chunks only, each a contiguous - # row-major (N, w_l) block; the backward needs only these derivatives. - saved = torch.cat([t.reshape(-1) for t in adots]) + # Chunk layout mirrors the CUDA op: the pre-activation of each layer as a + # contiguous row-major (N, w_l) block, before the bias. + saved = torch.cat([t.reshape(-1) for t in pres]) return e, saved @@ -210,30 +346,51 @@ def _cpu_backward( d_e: torch.Tensor, saved: torch.Tensor, ws: list[torch.Tensor], + bs: list[torch.Tensor], resnets: list[int], w_head: torch.Tensor, + act: int, ) -> torch.Tensor: total_width = sum(int(w.shape[1]) for w in ws) n_node = saved.shape[0] // total_width offset = [0] for w in ws: offset.append(offset[-1] + int(w.shape[1])) - adots = [ - saved[offset[l] * n_node : offset[l + 1] * n_node].reshape( - n_node, int(ws[l].shape[1]) - ) - for l in range(len(ws)) - ] dh = d_e.to(torch.float32) * w_head - for l in range(len(ws) - 1, -1, -1): - dpre = dh * adots[l] - dx = dpre @ ws[l].t() - if resnets[l] and ws[l].shape[0] == ws[l].shape[1]: + for layer in range(len(ws) - 1, -1, -1): + pre = saved[offset[layer] * n_node : offset[layer + 1] * n_node].reshape( + n_node, int(ws[layer].shape[1]) + ) + if bs[layer].numel(): + pre = pre + bs[layer] + dpre = dh * _activation_derivative(pre, act) + dx = dpre @ ws[layer].t() + if resnets[layer] and ws[layer].shape[0] == ws[layer].shape[1]: dx = dx + dh dh = dx return dh +def _cpu_energy_gradient( + x: torch.Tensor, + atype: torch.Tensor, + ws: list[torch.Tensor], + bs: list[torch.Tensor], + resnets: list[int], + w_head: torch.Tensor, + b_head: torch.Tensor, + bias_atom_e: torch.Tensor, + act: int, + seed: torch.Tensor, + tile: int, +) -> torch.Tensor: + energy, saved = _cpu_forward( + x, atype, ws, bs, resnets, w_head, b_head, bias_atom_e, act + ) + x.copy_(_cpu_backward(seed.reshape(-1, 1), saved, ws, bs, resnets, w_head, act)) + return energy + + # ====================================================================== # Registration and the public wrapper # ====================================================================== @@ -250,12 +407,70 @@ def ensure_registered() -> None: return torch.library.register_fake("deepmd::graph_fitting")(_forward_fake) torch.library.register_fake("deepmd::graph_fitting_backward")(_backward_fake) + torch.library.register_fake("deepmd::graph_fitting_energy_gradient")( + _energy_gradient_fake + ) torch.library.register_autograd( "deepmd::graph_fitting", _backward, setup_context=_setup_context ) _cpu_library = torch.library.Library("deepmd", "IMPL") _cpu_library.impl("graph_fitting", _cpu_forward, "CPU") _cpu_library.impl("graph_fitting_backward", _cpu_backward, "CPU") + _cpu_library.impl("graph_fitting_energy_gradient", _cpu_energy_gradient, "CPU") + + +def energy_and_input_gradient( + fit: Any, + descriptor: torch.Tensor, + atype: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-node energy and descriptor cotangent of an inference step. + + Inference seeds the fitting backward with the ownership mask, which is + known before the forward runs, so both directions are evaluated together + over node tiles and the pre-activations never reach node scale. The + cotangent is returned in the descriptor's own storage, which inference no + longer needs once the forward has consumed it. + + Parameters + ---------- + fit : EnergyFittingNet + The pt_expt fitting module. + descriptor : torch.Tensor + Flat descriptor with shape (N, nd), fp32 and contiguous. + atype : torch.Tensor + Flat node atom types with shape (N,), int64. + ownership : torch.Tensor + Mask selecting energy-contributing nodes with shape (N,). + atom_bias : torch.Tensor + Combined atomic energy bias with shape (ntypes,). + + Returns + ------- + atom_energy : torch.Tensor + Per-node energy with shape (N, 1), fp64, before the ownership mask. + descriptor_gradient : torch.Tensor + Cotangent of the descriptor with shape (N, nd), fp32. This is the + descriptor tensor itself, overwritten in place. + """ + ensure_registered() + network = fitting_operator_arguments(fit) + energy = torch.ops.deepmd.graph_fitting_energy_gradient( + descriptor, + atype, + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.head_bias, + atom_bias.to(torch.float64).contiguous(), + network.activation, + ownership.to(torch.float64).reshape(-1).contiguous(), + node_tile(), + ) + return energy, descriptor def graph_fitting( @@ -285,25 +500,16 @@ def graph_fitting( ``{fit.var_name: energy}`` with energy shape (N, 1), fp64. """ ensure_registered() - *hidden, head = fit.nets[0].layers - empty = hidden[0].w.new_empty(0) + arguments = fitting_operator_arguments(fit) e, _saved = torch.ops.deepmd.graph_fitting( descriptor.to(torch.float32).contiguous(), atype.contiguous(), - [layer.w.contiguous() for layer in hidden], - [layer.b.contiguous() if layer.b is not None else empty for layer in hidden], - [ - layer.idt.contiguous() if layer.idt is not None else empty - for layer in hidden - ], - [1 if layer.resnet else 0 for layer in hidden], - head.w.reshape(-1).contiguous(), - ( - head.b.reshape(-1).to(torch.float32).contiguous() - if head.b is not None - else empty - ), + arguments.weights, + arguments.biases, + arguments.residuals, + arguments.head_weight, + arguments.head_bias, fit.bias_atom_e.to(torch.float64).reshape(-1, 1)[:, 0].contiguous(), - ACT_CODES[str(hidden[0].activation_function).lower()], + arguments.activation, ) return {fit.var_name: e} diff --git a/deepmd/pt_expt/descriptor/__init__.py b/deepmd/pt_expt/descriptor/__init__.py index af1fa69893..0dc8847b5f 100644 --- a/deepmd/pt_expt/descriptor/__init__.py +++ b/deepmd/pt_expt/descriptor/__init__.py @@ -24,6 +24,9 @@ from .dpa4 import ( DescrptDPA4, ) +from .dpa4c import ( + DescrptDPA4C, +) from .hybrid import ( DescrptHybrid, ) @@ -49,6 +52,7 @@ "DescrptDPA2", "DescrptDPA3", "DescrptDPA4", + "DescrptDPA4C", "DescrptHybrid", "DescrptSeA", "DescrptSeAttenV2", diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py new file mode 100644 index 0000000000..21b3048c06 --- /dev/null +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""PyTorch-exportable execution backend for DPA4C. + +``deepmd.dpmodel.descriptor.dpa4c`` defines the graph algorithm and tensor +contracts. This module implements its performance-critical primitives with +native PyTorch operations, promotes DPA4C trainable arrays to parameters, +owns the mixed-precision and compressed-inference policies, and registers the +descriptor with the pt_expt backend. +""" + +from typing import ( + Any, + ClassVar, +) + +import torch + +from deepmd.dpmodel.descriptor.dpa4c import DescrptDPA4C as DescrptDPA4CDP +from deepmd.kernels.utils import ( + cuda_infer_level, + use_amp_infer, +) +from deepmd.pt_expt.common import ( + torch_module, +) +from deepmd.pt_expt.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.pt_expt.utils.update_sel import ( + UpdateSel, +) + +#: Learned arrays that the dpmodel keeps as buffers, keyed by owning class. +_TRAINABLE_ATTRS: dict[str, tuple[str, ...]] = { + "SeZMTypeEmbedding": ("adam_type_embedding",), + "RadialBasis": ("adam_freqs",), +} + + +def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module: + """Expose the dpmodel's learned buffers as PyTorch parameters. + + The two passes cannot be merged. Freezing recurses into children, so a + frozen parent can only detach a trainable child's array once that array + has become a parameter, which the first pass guarantees for the whole + tree before the second pass runs. + + Parameters + ---------- + module + Descriptor tree to promote in place. + + Returns + ------- + torch.nn.Module + The same module, for use as an expression. + """ + for submodule in module.modules(): + if not getattr(submodule, "trainable", True): + continue + for name in _TRAINABLE_ATTRS.get(type(submodule).__name__, ()): + value = submodule._buffers.get(name) + if value is None or not value.is_floating_point(): + continue + del submodule._buffers[name] + setattr(submodule, name, torch.nn.Parameter(value, requires_grad=True)) + + for submodule in module.modules(): + if not getattr(submodule, "trainable", True): + for parameter in submodule.parameters(recurse=True): + parameter.requires_grad_(False) + return module + + +@BaseDescriptor.register("dpa4c") +@torch_module +class DescrptDPA4C(DescrptDPA4CDP): + """Execute the backend-neutral DPA4C equations with PyTorch tensors. + + Notes + ----- + DPA4C components such as ``SeZMTypeEmbedding`` and ``RadialBasis`` store + trainable arrays in the dpmodel representation. The wrapper promotes those + arrays after construction and deserialization so they remain visible to + PyTorch optimizers and force-loss double backward. Graph gathers and cutoff + evaluation use native PyTorch primitives; the segment reductions and the + shared geometry and invariant modules already dispatch to PyTorch tensor + operations through the backend-neutral equations. + + This wrapper also owns the ``use_amp`` policy, because the array API has no + autocast and the dpmodel therefore only records the flag. + """ + + _update_sel_cls = UpdateSel + + #: Artifacts whose element type is not the ``float32`` the kernel consumes. + _COMPRESSION_BUFFER_DTYPES: ClassVar[dict[str, torch.dtype]] = { + "info": torch.float64, + "coupling_meta": torch.int32, + "coupling_entry": torch.int32, + } + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Construct and parameterize a PyTorch DPA4C descriptor. + + Parameters + ---------- + *args + Positional arguments forwarded to the dpmodel DPA4C constructor. + **kwargs + Keyword arguments forwarded to the dpmodel DPA4C constructor. + """ + super().__init__(*args, **kwargs) + _promote_trainable_tree(self) + self.compress = False + # Eval-time AMP is opted into through the environment and captured + # once, so a traced graph cannot depend on a later mutation. + self.use_amp_infer = use_amp_infer() + self._apply_autocast_policy() + + def call_graph( + self, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor | None = None, + comm_dict: dict | None = None, + ) -> tuple[torch.Tensor, None]: + """Evaluate the graph descriptor with compressed CUDA dispatch. + + Parameters + ---------- + graph + NeighborGraph over the flat node axis. + atype + Flat atom types with shape ``(N,)``. + type_embedding + Optional complete DPA4 type table. + comm_dict + Communication metadata accepted by the common graph ABI; unused. + + Returns + ------- + descriptor + Invariant descriptor with shape ``(N, get_dim_out())``. + rot_mat + ``None``. + """ + if ( + self.compress + and not self.training + and not self.exclude_types + and cuda_infer_level() >= 1 + and graph.destination_order is not None + and graph.destination_row_ptr is not None + ): + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + dpa4c_graph_compress, + mega_eligible, + op_available, + ) + + if op_available() and mega_eligible(self): + return dpa4c_graph_compress( + self, + graph, + atype, + ), None + if type_embedding is None: + type_embedding = self.type_embedding.call() + return super().call_graph( + graph, + atype, + type_embedding=type_embedding, + comm_dict=comm_dict, + ) + + def build_edge_features( + self, + graph: Any, + *args: Any, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the edge features under the DPA4C mixed-precision policy. + + The per-edge stage is the only region DPA4C autocasts. It holds every + tensor that scales with the edge count and every large matrix product, + namely the radial network and the pair-conditioned mode mixing, so + bfloat16 halves the dominant activation footprint. Autocast leaves the + geometry, the cutoff envelope, and the harmonics in full precision + because they are elementwise. + + The region ends at the returned features. The destination reduction + accumulates over the whole neighborhood and the readout raises the + moments to the fourth power, so both stay in the descriptor compute + precision; the ordered pair cache is likewise evaluated outside, over + the finite type table. + + Training follows ``use_amp`` and evaluation follows ``DP_AMP_INFER``. + The two are independent: mixed precision at inference is a throughput + choice that must not require a model to have been trained with it. + + Parameters + ---------- + graph + Neighbor graph in descriptor compute precision. + *args + Node types and ordered pair tables forwarded unchanged to the + backend-neutral implementation. + + Returns + ------- + amplitude + Masked edge amplitudes with shape ``(E, channels)``. + basis + Masked Cartesian harmonics with shape ``(E, (lmax + 1) ** 2)``. + envelope + Masked C³ envelope with shape ``(E,)``. + """ + autocast = graph.edge_vec.device.type == "cuda" and ( + self.use_amp if self.training else self.use_amp_infer + ) + if not autocast: + return super().build_edge_features(graph, *args) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): + features = super().build_edge_features(graph, *args) + dtype = graph.edge_vec.dtype + return tuple(feature.to(dtype) for feature in features) + + def _apply_autocast_policy(self) -> None: + """Let the layers inside the autocast region emit reduced precision. + + ``NativeLayer`` restores the dtype of its input after the affine map, + which would undo autocast at every layer of the radial network. Only + the radial trunk and the mode head sit inside the region, so only they + are opted out of that restoration, and only when mixed precision can + actually engage. A descriptor with neither switch set keeps the + default behavior exactly. + """ + enabled = self.use_amp or self.use_amp_infer + for layer in self.radial_embedding.layers: + layer.autocast_output = enabled + if self.radial_mode_head is not None: + self.radial_mode_head.autocast_output = enabled + + # === Backend primitives === + + def gather_rows( + self, + values: torch.Tensor, + index: torch.Tensor, + xp: Any | None = None, + ) -> torch.Tensor: + """Gather rows along the leading axis. + + Parameters + ---------- + values + Source tensor with shape ``(N, ...)``. + index + Row indices with shape ``(M,)``. Every DPA4C gather addresses a + flat node or edge axis, so a one-dimensional index suffices and + ``index_select`` avoids the advanced-indexing path. + xp + Array namespace accepted by the backend-neutral hook; unused. + + Returns + ------- + torch.Tensor + Gathered values with shape ``(M, *values.shape[1:])``. + """ + del xp + return torch.index_select(values, 0, index) + + def evaluate_cutoff_envelope(self, distance: torch.Tensor) -> torch.Tensor: + """Evaluate the fixed C³ cutoff envelope. + + Parameters + ---------- + distance + Regularized edge distances with shape ``(E, 1)`` in Å. + + Returns + ------- + torch.Tensor + Envelope values with shape ``(E, 1)``. + """ + u = torch.clamp( + (self.rcut - distance) / self.rcut, + min=0.0, + max=1.0, + ) + x = 1.0 - u + series = 1.0 + x * (4.0 + x * (10.0 + x * (20.0 + x * 35.0))) + return u**4 * series + + # === Statistics and parameter sharing === + + def compute_input_stats( + self, + merged: Any, + path: Any | None = None, + ) -> None: + """Calibrate output features without retaining an autograd graph.""" + with torch.no_grad(): + super().compute_input_stats(merged, path) + + def share_params( + self, + base_class: Any, + shared_level: int, + model_prob: float = 1.0, + resume: bool = False, + ) -> None: + """Reject compressed snapshots, then share as the base descriptor does. + + Assigning a submodule of ``base_class`` registers it in this module's + submodule table, so the backend-neutral implementation already + produces correct PyTorch sharing. Rebinding the whole submodule table + instead would additionally capture the pair-exclusion mask, which is + branch-local state. + + Parameters + ---------- + base_class + DPA4C descriptor that owns the shared parameters. + shared_level + Sharing level. Only level zero is supported. + model_prob + Model sampling probability accepted by the common ABI; unused. + resume + Checkpoint-restoration flag accepted by the common ABI; unused. + + Raises + ------ + RuntimeError + If either descriptor is a compressed inference snapshot. + """ + if self.compress or bool(getattr(base_class, "compress", False)): + raise RuntimeError( + "Compressed DPA4C snapshots cannot participate in parameter sharing." + ) + super().share_params(base_class, shared_level, model_prob, resume) + + # === Compressed-inference artifacts === + + @classmethod + def deserialize(cls, data: dict) -> "DescrptDPA4C": + """Deserialize DPA4C and restore trainable PyTorch parameters. + + Parameters + ---------- + data + Versioned dpmodel descriptor dictionary. + + Returns + ------- + DescrptDPA4C + Reconstructed PyTorch descriptor. + """ + data = data.copy() + compression = data.pop("compress", None) + obj = super().deserialize(data) + obj = _promote_trainable_tree(obj) + # Deserialization rebuilds the radial modules, so the autocast policy + # has to be reapplied to the fresh layers. + obj._apply_autocast_policy() + obj.compress = False + if compression is not None: + obj._set_compression( + { + name: torch.as_tensor(compression["@variables"][name]) + for name in obj._COMPRESSION_BUFFER_NAMES + } + ) + return obj + + def _set_compression( + self, + artifacts: dict[str, torch.Tensor], + ) -> None: + """Store immutable compressed-inference artifacts as module buffers. + + The metadata block keeps ``float64`` so that the radial table stride + and the cutoff survive without rounding, and the angular coupling + layout keeps ``int32``; every other artifact is the ``float32`` the + kernel consumes. + """ + device = next(self.parameters()).device + info = torch.as_tensor(artifacts["info"]) + self._compression_scalars = tuple( + float(value) for value in info.detach().cpu().tolist() + ) + for name in self._COMPRESSION_BUFFER_NAMES: + dtype = self._COMPRESSION_BUFFER_DTYPES.get(name, torch.float32) + value = artifacts[name].to(device=device, dtype=dtype).contiguous() + buffer_name = f"compress_{name}" + if buffer_name in self._buffers: + self._buffers[buffer_name] = value + else: + self.register_buffer(buffer_name, value) + self.compress = True + + def set_stat_mean_and_stddev(self, mean: Any, stddev: Any) -> None: + """Update output calibration and its compressed snapshot.""" + super().set_stat_mean_and_stddev(mean, stddev) + if self.compress: + device = self.compress_output_mean.device + self._buffers["compress_output_mean"] = torch.as_tensor( + mean, + dtype=torch.float32, + device=device, + ).contiguous() + self._buffers["compress_output_inv_std"] = torch.reciprocal( + torch.as_tensor( + stddev, + dtype=torch.float32, + device=device, + ) + ).contiguous() + + def train(self, mode: bool = True) -> "DescrptDPA4C": + """Set training mode while preserving compression immutability.""" + if mode and self.compress: + raise RuntimeError( + "A compressed DPA4C descriptor is an immutable inference " + "snapshot and cannot re-enter training mode." + ) + return super().train(mode) + + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 1.0, + table_stride_1: float = 0.002, + table_stride_2: float = 0.002, + check_frequency: int = -1, + ) -> None: + """Build immutable artifacts for the current DPA4C mega kernel. + + Parameters + ---------- + min_nbor_dist + Minimum neighbor distance accepted by the common compression ABI. + DPA4C tabulates the finite DPA4 radial basis from zero to ``rcut`` + and therefore does not use this value. + table_extrapolate + Common compression parameter; unused because the C³ radial map is + exactly zero beyond ``rcut``. + table_stride_1 + Uniform radial spline spacing in Å. + table_stride_2 + Common two-region table spacing; unused by the uniform DPA4C table. + check_frequency + Common overflow-check setting; unused because the radial domain is + bounded analytically. + + Raises + ------ + ValueError + If compression is already enabled or the descriptor configuration + has no compiled CUDA specialization. + """ + del min_nbor_dist, table_extrapolate, table_stride_2, check_frequency + if self.compress: + raise ValueError("Compression is already enabled.") + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + build_compression_artifacts, + ) + + self._set_compression(build_compression_artifacts(self, table_stride_1)) + + def fused_energy_force_graph( + self, + fitting: Any, + graph: Any, + atype: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + do_atomic_virial: bool, + ) -> ( + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ] + | None + ): + """Evaluate the inference-only compressed energy-force composition. + + Returns ``None`` when the model or graph cannot use the level-two CUDA + path, allowing the caller to retain the generic autograd lower. + """ + if ( + self.training + or not self.compress + or bool(self.exclude_types) + or cuda_infer_level() < 2 + or graph.destination_order is None + or graph.destination_row_ptr is None + or graph.source_order is None + or graph.source_row_ptr is None + ): + return None + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + dpa4c_graph_compress_energy_force, + ef_op_available, + mega_eligible, + ) + from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + ) + + if ( + not ef_op_available() + or not mega_eligible(self) + or not fitting_eligible(fitting) + ): + return None + return dpa4c_graph_compress_energy_force( + self, + fitting, + graph, + atype, + ownership, + atom_bias, + atype.shape[0], + do_atomic_virial, + ) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 66ddd1de79..5a5f8d4e06 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -357,7 +357,7 @@ def update_neighbor_stat( options: TrainEntrypointOptions, *, multi_task: bool, - ) -> tuple[dict[str, Any], None]: + ) -> tuple[dict[str, Any], float | dict[str, float | None] | None]: """Update pt_expt descriptor selections from neighbor statistics.""" log.info( "Calculate neighbor statistics... " @@ -367,19 +367,23 @@ def update_neighbor_stat( BaseModel, ) + min_nbor_dist: dict[str, float | None] = {} for task_config in iter_training_task_configs(config): type_map = task_config.model_params.get("type_map") train_data = _get_neighbor_stat_data( dict(task_config.training_data_params), type_map ) - updated_model_params, _ = BaseModel.update_sel( + updated_model_params, task_min_nbor_dist = BaseModel.update_sel( train_data, type_map, dict(task_config.model_params) ) + min_nbor_dist[task_config.key] = task_min_nbor_dist if multi_task: config["model"]["model_dict"][task_config.key] = updated_model_params else: config["model"] = updated_model_params - return config, None + if multi_task: + return config, min_nbor_dist + return config, next(iter(min_nbor_dist.values()), None) def print_summary(self) -> None: """Print pt_expt backend summary.""" @@ -429,6 +433,13 @@ def run_training( finetune_links=self.finetune_links, shared_links=self.shared_links, ) + # Persist the neighbor statistic on the model so that `dp compress` + # reads it from the checkpoint instead of rescanning the training data. + if isinstance(neighbor_stat, dict): + for task_key, task_min_nbor_dist in neighbor_stat.items(): + trainer.model[task_key].min_nbor_dist = task_min_nbor_dist + elif neighbor_stat is not None: + trainer.model.min_nbor_dist = neighbor_stat trainer.run() diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index b084f3fdef..db32f0ceec 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -247,7 +247,11 @@ def __init__( # where the corresponding builder knob is ``nlist_backend``. if neighbor_graph_method != "auto" and getattr(self, "metadata", {}).get( "lower_input_kind" - ) not in ("graph", "dpa1_canonical"): + ) not in ( + "graph", + "dpa1_canonical", + "dpa4c_canonical", + ): raise ValueError( f"neighbor_graph_method={neighbor_graph_method!r} only applies to " "graph-routed artifacts (lower_input_kind == 'graph'); this " @@ -312,7 +316,11 @@ def _setup_neighbor_backend(self, nlist_backend: str) -> None: f"Unknown nlist_backend '{nlist_backend}'; " "expected 'auto', 'vesin', or 'native'." ) - if self.metadata.get("lower_input_kind") in ("graph", "dpa1_canonical"): + if self.metadata.get("lower_input_kind") in ( + "graph", + "dpa1_canonical", + "dpa4c_canonical", + ): if self.neighbor_list is not None: raise ValueError( "neighbor_list cannot be used with this graph-routed model: " @@ -1753,7 +1761,11 @@ def _eval_model( request_defs: list[OutputVariableDef], charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: - if self.metadata.get("lower_input_kind") in ("graph", "dpa1_canonical"): + if self.metadata.get("lower_input_kind") in ( + "graph", + "dpa1_canonical", + "dpa4c_canonical", + ): return self._eval_model_graph( coords, cells, atom_types, fparam, aparam, request_defs, charge_spin ) @@ -2204,7 +2216,10 @@ def _eval_model_graph( device=DEVICE, ) - if self.metadata.get("lower_input_kind") == "dpa1_canonical": + if self.metadata.get("lower_input_kind") in ( + "dpa1_canonical", + "dpa4c_canonical", + ): # The canonical ABI has NO fparam/aparam/charge_spin slots; the # export gate (fitting_eligible) rejects such models today, so # this is unreachable -- assert it loudly so a future loosening @@ -2216,10 +2231,10 @@ def _eval_model_graph( or int(self.metadata.get("dim_chg_spin", 0) or 0) > 0 ): raise NotImplementedError( - "dpa1_canonical artifacts carry no fparam/aparam/" + "compact canonical artifacts carry no fparam/aparam/" "charge_spin inputs; a model requiring them must not be " - "frozen with lower_kind='dpa1_canonical' (the export " - "eligibility gate should have rejected it)." + "frozen with a canonical lower kind (the export eligibility " + "gate should have rejected it)." ) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, @@ -2566,6 +2581,10 @@ def eval_descriptor( Frame parameters, optional. aparam Atom parameters, optional. + charge_spin + Optional frame-level charge and spin conditioning. + **kwargs + Additional backend-compatible evaluation options. Returns ------- @@ -2635,6 +2654,10 @@ def eval_fitting_last_layer( Frame parameters, optional. aparam Atom parameters, optional. + charge_spin + Optional frame-level charge and spin conditioning. + **kwargs + Additional backend-compatible evaluation options. Returns ------- diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 6aa5f4add4..90c85cfd5c 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -67,7 +67,7 @@ def forward_lower_canonical_graph( *, do_atomic_virial: bool, ) -> dict[str, torch.Tensor]: - """Evaluate an eligible compressed DPA1 deployment graph. + """Evaluate an eligible compressed canonical deployment graph. Parameters ---------- @@ -78,7 +78,7 @@ def forward_lower_canonical_graph( n_local Per-frame owned node counts with shape ``(nf,)``, int64. source - Source-node indices with shape ``(S,)``, int32 or int64. + Source-node indices with shape ``(S,)``, uint32. edge_vec Destination-major edge vectors with shape ``(S, 3)``, float32. destination_row_ptr @@ -97,17 +97,26 @@ def forward_lower_canonical_graph( Public energy-model outputs on the flat node axis. """ from deepmd.kernels.cuda.dpa1.canonical import ( - canonical_model_eligible, + canonical_model_eligible as dpa1_canonical_eligible, + ) + from deepmd.kernels.cuda.dpa1.canonical import ( dpa1_canonical_compress_energy_force, ) + from deepmd.kernels.cuda.dpa4c.canonical import ( + canonical_model_eligible as dpa4c_canonical_eligible, + ) + from deepmd.kernels.cuda.dpa4c.canonical import ( + dpa4c_canonical_compress_energy_force, + ) from deepmd.pt_expt.utils.canonical_graph import ( - DPA1CanonicalGraph, + CanonicalGraph, validate_canonical_graph_shapes, ) - if not canonical_model_eligible(self): + use_dpa4c = dpa4c_canonical_eligible(self) + if not use_dpa4c and not dpa1_canonical_eligible(self): raise ValueError("model is not eligible for compact canonical deployment") - graph = DPA1CanonicalGraph( + graph = CanonicalGraph( n_node=n_node, n_local=n_local, source=source, @@ -126,21 +135,33 @@ def forward_lower_canonical_graph( descriptor = self.atomic_model.descriptor fitting = self.atomic_model.fitting_net atom_bias = fitting.bias_atom_e[:, 0] + self.atomic_model.out_bias[0, :, 0] - energy, atom_energy, force, virial, atom_virial = ( - dpa1_canonical_compress_energy_force( - descriptor, - fitting, - graph, - atype, - # descriptor-owned hook (single owner for the graph-route tebd - # table); value-identical for dpa1, the only canonical-eligible - # descriptor. - descriptor.graph_type_embedding_table(), - output_mask, - atom_bias, - do_atomic_virial, + if use_dpa4c: + energy, atom_energy, force, virial, atom_virial = ( + dpa4c_canonical_compress_energy_force( + descriptor, + fitting, + graph, + atype, + output_mask, + atom_bias, + do_atomic_virial, + ) + ) + else: + energy, atom_energy, force, virial, atom_virial = ( + dpa1_canonical_compress_energy_force( + descriptor, + fitting, + graph, + atype, + # Descriptor-owned hook: the single owner of the + # graph-route type-embedding table. + descriptor.graph_type_embedding_table(), + output_mask, + atom_bias, + do_atomic_virial, + ) ) - ) result = { "atom_energy": atom_energy, "energy": energy, @@ -236,6 +257,20 @@ def forward( Parameters ---------- + coord + Atomic coordinates. + atype + Atomic type indices. + box + Simulation-cell vectors, or ``None`` for a non-periodic system. + fparam + Optional frame parameters. + aparam + Optional atomic parameters. + do_atomic_virial + Whether to return per-atom virials. + charge_spin + Optional frame-level charge and spin conditioning. neighbor_list The neighbor-list construction strategy forwarded to :meth:`call_common`. ``None`` uses the default all-pairs builder @@ -425,8 +460,22 @@ def forward_lower_exportable( Parameters ---------- - extended_coord, extended_atype, nlist, mapping, fparam, aparam, do_atomic_virial - Sample inputs with representative shapes (used for tracing). + extended_coord + Extended-coordinate sample used for tracing. + extended_atype + Extended atom-type sample used for tracing. + nlist + Neighbor-list sample used for tracing. + mapping + Extended-to-local mapping sample used for tracing. + fparam + Optional frame-parameter sample. + aparam + Optional atomic-parameter sample. + do_atomic_virial + Whether the traced module returns per-atom virials. + charge_spin + Optional charge/spin conditioning sample. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -686,6 +735,18 @@ def forward_lower_graph_exportable_with_comm( atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin, do_atomic_virial As in :meth:`forward_lower_graph_exportable`. + destination_order + Destination-major edge permutation used by fused graph operators. + + destination_row_ptr + Destination CSR row pointers. + + source_order + Source-major edge permutation used by force assembly. + + source_row_ptr + Source CSR row pointers. + send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost The 8 comm tensors (see ``_make_comm_sample_inputs`` in ``serialization.py``), packed into ``comm_dict`` inside the diff --git a/deepmd/pt_expt/model/graph_lower.py b/deepmd/pt_expt/model/graph_lower.py index f79c27a6ca..5fa111661a 100644 --- a/deepmd/pt_expt/model/graph_lower.py +++ b/deepmd/pt_expt/model/graph_lower.py @@ -25,10 +25,10 @@ def graph_edge_dtype(model: Any, lower_kind: str) -> str: Returns ------- str - ``"float32"`` for eligible geometrically compressed DPA1 graph - lowers, otherwise ``"float64"``. + ``"float32"`` for eligible compressed DPA1 or DPA4C graph lowers, + otherwise ``"float64"``. """ - if lower_kind not in ("graph", "dpa1_canonical"): + if lower_kind not in ("graph", "dpa1_canonical", "dpa4c_canonical"): return "float64" return str(model.atomic_model.graph_edge_dtype()) diff --git a/deepmd/pt_expt/utils/canonical_graph.py b/deepmd/pt_expt/utils/canonical_graph.py index 5db225ab1a..e01163a073 100644 --- a/deepmd/pt_expt/utils/canonical_graph.py +++ b/deepmd/pt_expt/utils/canonical_graph.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Compact canonical graph contract for compressed DPA1 deployment.""" +"""Compact canonical graph contract for compressed CUDA deployment.""" from __future__ import ( annotations, @@ -19,9 +19,11 @@ NeighborGraph, ) +UINT32_MAX = (1 << 32) - 1 + @dataclass(frozen=True) -class DPA1CanonicalGraph: +class CanonicalGraph: """Store a cutoff-compacted destination-major graph without redundant fields. The physical edge count is the final value of either CSR row-pointer tensor. @@ -36,7 +38,7 @@ class DPA1CanonicalGraph: Per-frame owned node counts with shape ``(nf,)``, int64. source Source-node index for each edge storage slot with shape ``(S,)``, - int64. + uint32. edge_vec Neighbor-minus-center vectors with shape ``(S, 3)``, float32. destination_row_ptr @@ -58,14 +60,14 @@ class DPA1CanonicalGraph: def validate_canonical_graph_shapes( - graph: DPA1CanonicalGraph, + graph: CanonicalGraph, node_count: int, ) -> None: """Validate shape, dtype, and device invariants without reading tensor data.""" index_dtype = graph.source.dtype - if index_dtype != torch.int64: - raise ValueError("canonical graph source must be int64") - if graph.source_order.dtype != torch.int64: + if index_dtype != torch.uint32: + raise ValueError("canonical graph source must be uint32") + if graph.source_order.dtype != index_dtype: raise ValueError("canonical graph source and source_order dtypes must match") if graph.edge_vec.dtype != torch.float32: raise ValueError("canonical graph edge_vec must be float32") @@ -82,6 +84,8 @@ def validate_canonical_graph_shapes( raise ValueError("canonical graph edge_vec must have shape (S, 3)") if graph.source.shape[0] < 2: raise ValueError("canonical graph edge storage must contain at least two slots") + if graph.source.shape[0] > UINT32_MAX: + raise ValueError("canonical graph edge storage exceeds the uint32 range") if graph.destination_row_ptr.shape != (node_count + 1,): raise ValueError("destination_row_ptr must have shape (N + 1,)") if graph.source_row_ptr.shape != (node_count + 1,): @@ -114,7 +118,7 @@ def validate_canonical_graph_shapes( def canonical_graph_from_neighbor_graph( graph: NeighborGraph, -) -> DPA1CanonicalGraph: +) -> CanonicalGraph: """Convert a compact generic graph into the source-only deployment contract. Parameters @@ -124,7 +128,7 @@ def canonical_graph_from_neighbor_graph( Returns ------- - DPA1CanonicalGraph + CanonicalGraph Source-only graph with exactly two storage slots when ``E < 2``. Raises @@ -156,9 +160,11 @@ def canonical_graph_from_neighbor_graph( node_count = graph.destination_row_ptr.shape[0] - 1 storage_edge_count = max(physical_edge_count, 2) + if storage_edge_count > UINT32_MAX: + raise ValueError("canonical graph edge storage exceeds the uint32 range") source = torch.zeros( storage_edge_count, - dtype=torch.int64, + dtype=torch.uint32, device=graph.edge_index.device, ) edge_vec = torch.zeros( @@ -171,19 +177,19 @@ def canonical_graph_from_neighbor_graph( storage_edge_count, dtype=torch.int64, device=graph.edge_index.device, - ) + ).to(torch.uint32) if physical_edge_count: source[:physical_edge_count] = graph.edge_index[0, :physical_edge_count].to( - torch.int64 + torch.uint32 ) edge_vec[:physical_edge_count] = graph.edge_vec[:physical_edge_count].to( torch.float32 ) source_order[:physical_edge_count] = graph.source_order[ :physical_edge_count - ].to(torch.int64) + ].to(torch.uint32) - result = DPA1CanonicalGraph( + result = CanonicalGraph( n_node=graph.n_node.contiguous(), n_local=graph.n_local.contiguous(), source=source, @@ -197,7 +203,8 @@ def canonical_graph_from_neighbor_graph( __all__ = [ - "DPA1CanonicalGraph", + "UINT32_MAX", + "CanonicalGraph", "canonical_graph_from_neighbor_graph", "validate_canonical_graph_shapes", ] diff --git a/deepmd/pt_expt/utils/network.py b/deepmd/pt_expt/utils/network.py index f3ca3b392d..adc7d4f326 100644 --- a/deepmd/pt_expt/utils/network.py +++ b/deepmd/pt_expt/utils/network.py @@ -81,6 +81,12 @@ class NativeLayer(NativeLayerDP, torch.nn.Module): see the ``TorchArrayParam`` docstring. """ + # Restoring the dtype of the input after the affine map would undo an + # enclosing ``torch.autocast`` region and defeat mixed precision across a + # stack of these layers. An owner that autocasts sets this to ``True`` on + # the layers inside its region; every other layer keeps the default. + autocast_output: bool = False + def __init__(self, *args: Any, **kwargs: Any) -> None: torch.nn.Module.__init__(self) NativeLayerDP.__init__(self, *args, **kwargs) @@ -144,7 +150,7 @@ def call(self, x: torch.Tensor) -> torch.Tensor: if self.b is not None else torch.matmul(x, self.w) ) - if y.dtype != x.dtype: + if not self.autocast_output and y.dtype != x.dtype: y = y.to(x.dtype) y = _torch_activation(y, self.activation_function) if self.idt is not None: diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 5d7be558bb..4abc487ad3 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -615,7 +615,7 @@ def build_synthetic_canonical_graph_inputs( *, device: torch.device, ) -> tuple[torch.Tensor, ...]: - """Build the compact canonical trace inputs for compressed DPA1.""" + """Build compact canonical trace inputs for compressed CUDA descriptors.""" from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) @@ -678,9 +678,17 @@ def _build_canonical_graph_dynamic_shapes( ) -> tuple: """Build dynamic shapes for the eight-tensor compact deployment ABI.""" del sample_inputs + from deepmd.pt_expt.utils.canonical_graph import ( + UINT32_MAX, + ) + nframes_dim = torch.export.Dim("nframes", min=1) node_dim = torch.export.Dim("n_node_total", min=1) - edge_storage_dim = torch.export.Dim("nedge_storage", min=2) + edge_storage_dim = torch.export.Dim( + "nedge_storage", + min=2, + max=UINT32_MAX, + ) return ( {0: node_dim}, {0: nframes_dim}, @@ -922,8 +930,13 @@ def _build_dynamic_shapes( Whether the inputs include the 8 comm tensors. model_nnei : int The model's sum(sel). Used as the min for the dynamic nnei dim. - Returns a tuple (not dict) to match positional args of the make_fx - traced module, whose arg names may have suffixes like ``_1``. + + Returns + ------- + tuple + Dynamic-shape specifications in the positional order of the make_fx + traced module. A tuple is required because traced argument names may + carry generated suffixes such as ``_1``. """ # When tracing the with-comm variant, nframes is static at 1. # Rationale: pt_expt's Repflow/Repformer parallel-mode override @@ -1123,6 +1136,8 @@ def _probe_has_message_passing(obj: object) -> bool | None: # The C++ loader branches on this to build the matching inputs. meta["lower_input_kind"] = lower_kind meta["graph_edge_dtype"] = graph_edge_dtype(model, lower_kind) + if lower_kind in ("dpa1_canonical", "dpa4c_canonical"): + meta["canonical_index_dtype"] = "uint32" # Model-level pair-type exclusion (``pair_exclude_types``): a list of # ``[ti, tj]`` type pairs whose interaction is dropped. Exclusion is a @@ -1250,7 +1265,8 @@ def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str: ``"auto"`` selects the graph lower for a graph-lower model whose graph implementation is exportable to ``.pt2`` and the dense nlist lower for - everything else. An explicit ``"nlist"`` / ``"graph"`` is returned + everything else. Eligible compressed DPA1 and DPA4C energy models select + their compact canonical graph schemas. Any explicit lower kind is returned unchanged. """ if lower_kind != "auto": @@ -1267,10 +1283,17 @@ def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str: model = BaseModel.deserialize(data["model"]) if model_uses_graph_lower(model) and _supports_graph_export(model): from deepmd.kernels.cuda.dpa1.canonical import ( - canonical_model_eligible, + canonical_model_eligible as dpa1_canonical_eligible, + ) + from deepmd.kernels.cuda.dpa4c.canonical import ( + canonical_model_eligible as dpa4c_canonical_eligible, ) - return "dpa1_canonical" if canonical_model_eligible(model) else "graph" + if dpa4c_canonical_eligible(model): + return "dpa4c_canonical" + if dpa1_canonical_eligible(model): + return "dpa1_canonical" + return "graph" return "nlist" @@ -1336,7 +1359,7 @@ def deserialize_to_file( # DP_CUDA_INFER >= 2 so the analytic backward and CSR scatter remain custom # operators, while the per-atom virial is mandatory for the LAMMPS Kokkos # consumer. - if lower_kind in ("graph", "dpa1_canonical"): + if lower_kind in ("graph", "dpa1_canonical", "dpa4c_canonical"): do_atomic_virial = True ctx: contextlib.AbstractContextManager = _cuda_infer_at_least_2() else: @@ -1463,7 +1486,7 @@ def _trace_and_export( ) # Graph-form exports use a dynamic edge axis and an energy-model contract. - if lower_kind in ("graph", "dpa1_canonical"): + if lower_kind in ("graph", "dpa1_canonical", "dpa4c_canonical"): import math check_graph_trace_torch_version(model) @@ -1509,7 +1532,7 @@ def _trace_and_export( "forward_lower_graph_exportable_with_comm; graph-form " "with-comm .pt2 export requires an energy model" ) - canonical = lower_kind == "dpa1_canonical" + canonical = lower_kind in ("dpa1_canonical", "dpa4c_canonical") required_method = ( "forward_lower_canonical_graph_exportable" if canonical @@ -1520,14 +1543,19 @@ def _trace_and_export( f"model {type(model).__name__} has no {required_method}" ) if canonical: - from deepmd.kernels.cuda.dpa1.canonical import ( - canonical_model_eligible, - ) + if lower_kind == "dpa4c_canonical": + from deepmd.kernels.cuda.dpa4c.canonical import ( + canonical_model_eligible, + ) + else: + from deepmd.kernels.cuda.dpa1.canonical import ( + canonical_model_eligible, + ) if not canonical_model_eligible(model): raise NotImplementedError( "compact canonical export requires an eligible compressed " - "DPA1 energy model" + "DPA1 or DPA4C energy model" ) # Trace-time sizes must be pairwise-distinct AND avoid every static diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index d9d57dfb1a..d0f82ebe49 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -451,6 +451,121 @@ def descrpt_se_a_args() -> list[Argument]: ] +@descrpt_args_plugin.register( + "dpa4c", + alias=["DPA4C"], + doc=doc_only_pt_expt_supported + + "DPA4C is the compact and compressible degree-wise descriptor of the DPA4 family.", +) +def descrpt_dpa4c_args() -> list[Argument]: + """Return the DPA4C descriptor arguments.""" + return [ + Argument( + "rcut", + float, + optional=True, + default=6.0, + doc="The outer cutoff radius.", + ), + Argument( + "channels", + int, + optional=True, + default=32, + doc=( + "Scalar degree-zero and edge channel width. Supported values " + "are 8, 16, 32, 64, and 128. This is the primary scaling " + "knob: it widens the edge features, the per-atom angular " + "state, and the descriptor output together. The fitting " + "network is sized against it; the released grades pair " + "channels 8, 32, 64, and 128 with fitting hidden widths 96, " + "192, 256, and 384." + ), + ), + Argument( + "lmax", + int, + optional=True, + default=2, + doc="Maximum angular degree. Supported values are 2, 3, and 4.", + ), + Argument( + "basis_type", + str, + optional=True, + default="bessel", + doc="DPA4 radial basis type: `bessel` or `gaussian`.", + ), + Argument( + "n_radial", + int, + optional=True, + default=16, + doc=( + "Number of DPA4 radial basis functions forming the fixed " + "analytic radial input." + ), + ), + Argument( + "radial_modes", + int, + optional=True, + default=0, + doc=( + "Number of shared radial mode profiles that every ordered " + "atom-type pair mixes with its own coefficients. Zero leaves " + "each pair with a rescaled copy of one shared radial " + "function; larger values let each pair select its own radial " + "shape." + ), + ), + Argument( + "use_amp", + bool, + optional=True, + default=False, + doc=( + "If True, run the per-edge stage under bfloat16 automatic " + "mixed precision on CUDA during training. This lowers the " + "dominant activation footprint, which scales with the edge " + "count. The destination reduction and the invariant readout " + "stay in the descriptor precision. Evaluation and inference " + "are governed independently by the `DP_AMP_INFER` environment " + "variable, so a model trained in full precision can still " + "infer under mixed precision, and the reverse." + ), + ), + Argument( + "exclude_types", + list[list[int]], + optional=True, + default=[], + doc="Ordered atom-type pairs excluded from the descriptor.", + ), + Argument( + "precision", + str, + optional=True, + default="float32", + doc="Floating-point precision of descriptor parameters.", + ), + Argument( + "trainable", + bool, + optional=True, + default=True, + doc="Whether descriptor parameters are trainable.", + ), + Argument( + "seed", + [int, None], + optional=True, + default=None, + doc="Random seed for parameter initialization.", + ), + ] + + @descrpt_args_plugin.register( "dpa4", alias=["DPA4", "SeZM", "sezm"], diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md new file mode 100644 index 0000000000..c1a217f1af --- /dev/null +++ b/doc/model/dpa4c.md @@ -0,0 +1,334 @@ +# Descriptor DPA4C {{ pytorch_icon }} + +> [!NOTE] +> **Supported backends**: PyTorch Exportable {{ pytorch_icon }} (`dp --pt-expt`) + +DPA4C is the compact and compressible degree-wise descriptor of the DPA4 +family. Where DPA4/SeZM targets the accuracy frontier through equivariant +message passing, DPA4C targets the throughput frontier: it reads each local +environment once, keeps no message-passing state, and admits a compressed CUDA +inference path in which its radial functions are replaced by tabulated splines. +It is intended for large-scale molecular dynamics and as a distillation student +of a DPA4 teacher. + +DPA4C is selected as a descriptor, `descriptor.type: "dpa4c"`, and pairs with +the standard energy fitting network. There is no separate `model.type` scaffold. + +## Quick start + +```bash +cd examples/water/dpa4c +dp --pt-expt train input.json +``` + +`examples/water/dpa4c/input.json` is a complete energy-training input you can +copy and adapt. See [training energy models](train-energy.md) for the general +workflow shared by all energy models. + +## Overview + +DPA4C predicts atomic energies and obtains forces and virials by +differentiating the energy, the same conservative formulation used by every +standard DeePMD energy model: + +```math +\mathbf{F}_i = -\frac{\partial E}{\partial \mathbf{r}_i}. +``` + +For each atom the descriptor accumulates angular moments of its neighbors up to +degree {ref}`lmax `, contracts them into +rotationally invariant scalars, and passes only those scalars to the fitting +network. The neighbor shell is read exactly once: there is no message passing, +so an atom's descriptor depends only on the atoms within +{ref}`rcut ` of it. This one-hop +locality is what keeps the per-step cost low and makes the compressed inference +path possible. + +Two properties follow from the construction and matter in practice. The radial +map is exactly zero at and beyond `rcut`, with continuous derivatives, so the +potential energy surface stays smooth as neighbors cross the cutoff. And every +per-atom quantity is bounded analytically, which is why compression needs +neither an extrapolation region nor overflow checking. + +If you want the design details, see +[Architecture details](#architecture-details) at the end of this page. + +## Configuration + +### Minimal input + +A minimal DPA4C descriptor needs nothing but its type; every option has a +documented default. + +```json +{ + "model": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "dpa4c", + "rcut": 6.0 + }, + "fitting_net": { + "neuron": [ + 128, + 128, + 128 + ], + "activation_function": "silu" + } + } +} +``` + +DPA4C has no `sel` option. It is graph-native: the descriptor consumes a +carry-all neighbor graph that holds every neighbor within `rcut`, rather than a +fixed-capacity neighbor list. There is no capacity to size, no dependence on the +densest frame in the dataset, and no truncation to guard against. + +DPA4C defaults to `float32` +({ref}`precision `), which is also +what the compressed CUDA path requires. Double precision is neither necessary +nor supported for compressed inference. + +### Main options + +Every option, with its default and full description, is listed in the +{ref}`argument reference `. Four of them +carry the accuracy–cost trade-off: + +- **Width** — {ref}`channels `, one + of 8, 16, 32, 64, or 128. This is the primary scaling knob. It widens the + scalar and edge features, the per-atom angular state, and the descriptor + output together, so it costs both throughput and the largest system that fits + in memory. +- **Angular degree** — {ref}`lmax `, one + of 2, 3, or 4. Each additional degree adds angular components to the per-atom + state. Its absolute cost is fixed by the degree, so its *relative* cost is + largest at narrow widths. +- **Radial resolution** — + {ref}`radial_modes `. Zero + leaves every ordered atom-type pair with a rescaled copy of one shared radial + function; larger values let each pair select its own radial shape from several + shared profiles. It spends per-edge work without enlarging the per-atom state, + which makes it the lever to reach for when memory rather than throughput is + the binding constraint. +- **Radial basis** — + {ref}`basis_type ` and + {ref}`n_radial ` select the + analytic basis that feeds the radial network. + +> [!IMPORTANT] +> Compressed inference is compiled for +> `radial_modes` in `{0, 2, 4, 8}` only. A model trained with any other value +> trains and runs correctly on the portable path, but `dp --pt-expt compress` +> will reject it. Choose the value with compression in mind if you intend to +> deploy the compressed model. + +### Recommended configurations + +The released grades pair each descriptor width with a fitting width sized +against it, in ascending cost. They are good starting points; `Neo` is the +general-purpose default. + +| Grade | `channels` | `lmax` | `radial_modes` | Fitting hidden width | +| ----- | ---------: | -----: | -------------: | -------------------: | +| Nano | 8 | 2 | 0 | 96 | +| Mini | 32 | 2 | 0 | 192 | +| Neo | 32 | 2 | 4 | 192 | +| Air | 64 | 3 | 4 | 256 | +| Plus | 128 | 3 | 4 | 384 | + +`Mini` and `Neo` share a descriptor width and differ only in the radial modes, +which buys accuracy at a per-edge cost while leaving the largest tractable +system unchanged. `Air` and `Plus` additionally raise the angular degree. + +The fitting network is sized against the descriptor because the invariant +output grows with `channels`. Unlike `radial_modes`, fitting width is not a free +trade against memory: it adds per-atom activations and the derivatives saved for +the force backward pass, so widening or deepening it costs throughput and +capacity together. Widen it only when validation error is limited by fitting +capacity rather than by the descriptor. + +## Training + +The recommended objective is the standard conservative energy loss: + +```json +{ + "loss": { + "type": "ener" + } +} +``` + +See [training energy models](train-energy.md) for the general workflow. + +### Mixed precision + +{ref}`use_amp ` runs the per-edge +stage under bfloat16 automatic mixed precision on CUDA during training. The +activation footprint of that stage scales with the edge count and dominates +memory, so enabling it lowers peak memory substantially; the destination +reduction and the invariant readout stay in the descriptor precision. Use it on +GPUs with native bf16 support. + +`use_amp` is an execution policy rather than model state: it is not serialized, +and evaluation and inference are governed independently by `DP_AMP_INFER`. A +model trained in full precision can therefore be evaluated under mixed +precision, and the reverse. + +## Model compression + +Compression replaces the analytic radial functions and their type-pair +modulation with tabulated splines evaluated by fused CUDA kernels. Because the +radial map is analytically bounded and vanishes at `rcut`, the table needs no +extrapolation region and no overflow checking. + +The workflow is the standard three steps: + +```bash +dp --pt-expt train input.json +dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph +dp --pt-expt compress -i frozen_model.pt2 -o compressed_model.pt2 +``` + +Only `-s, --step` applies to DPA4C; it sets the uniform spline spacing in Å, and +a smaller value means a finer table and a larger model. +The `--extrapolate`, `--frequency`, and `--training-script` options exist for +descriptors whose tables need a second region, an overflow guard, or a minimum +neighbor distance computed from data; DPA4C needs none of them and ignores them. + +Compression requires: + +- the PyTorch Exportable backend on CUDA; +- `precision: "float32"`; +- `channels` in `{8, 16, 32, 64, 128}`, `lmax` in `{2, 3, 4}`, and + `radial_modes` in `{0, 2, 4, 8}`; +- an empty + {ref}`exclude_types `, since + the fused kernel has no type-exclusion branch. A compressed model with + excluded pairs falls back to the portable path. + +`dp --pt-expt compress` reports an explicit error when the configuration falls +outside these sets. + +## Export and running in LAMMPS + +DPA4C uses the PyTorch `.pt2` (AOTInductor) export path. Freeze with the graph +lower, which is the form the C++ graph path consumes: + +```bash +dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph +``` + +Use the frozen or compressed `.pt2` with the `deepmd` pair style: + +```lammps +pair_style deepmd compressed_model.pt2 +pair_coeff * * O H +``` + +Because DPA4C performs no message passing, it needs no cross-rank halo exchange +of intermediate features, and MPI domain decomposition follows the ordinary +pair-style path. Launch one MPI rank per GPU and make every target device +visible: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 mpirun -np 4 lmp -in in.lammps +``` + +Use a non-zero neighbor skin, for example `neighbor 2.0 bin`, to keep per-step +GPU memory stable; a zero skin rebuilds the neighbor list every step. + +## Inference settings + +Inference behavior is controlled by environment variables read when the model is +constructed: + +| Environment variable | Default | Effect | +| -------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `DP_CUDA_INFER` | `0` | Fused CUDA kernel level. `0` disables them. `1` uses the fused descriptor and fitting operators with the force from autograd. `2` additionally collapses descriptor, fitting, and force/virial assembly into one operator, numerically identical to level 1. | +| `DP_AMP_INFER` | off | bf16 autocast over the per-edge stage during inference. Independent of the training-time `use_amp`. | +| `DP_TF32_INFER` | `0` | float32 matmul precision: `0` highest, `1` high, `2` medium. | + +A compressed model requires `DP_CUDA_INFER` of at least `1` to reach its fused +path; at `0` it evaluates through the portable path and the compression brings +no speedup. For molecular dynamics sensitive to the smoothness of the potential +energy surface, keep `DP_TF32_INFER=0` and `DP_AMP_INFER=0`. + +> [!IMPORTANT] +> Set these variables **before** running `dp --pt-expt freeze` or +> `dp --pt-expt compress`. The exported `.pt2` is an AOTInductor artifact, so the +> kernel level and precision policy are captured into the graph at export time +> and are **not** re-evaluated when the `.pt2` is later loaded by LAMMPS. + +## Data format + +DPA4C consumes a mixed-type neighbor list, so it supports both the +[standard DeePMD-kit data format](../data/system.md) and the +[mixed-type data format](../data/system.md#mixed-type). Keep the `type_map` +order consistent across the dataset, the input file, and any downstream +`pair_coeff` mapping. + +## Architecture details + +Optional background on how the descriptor works, linking each part to the +options that control it. Skip it unless you are tuning those options. + +### Edge features + +For every neighbor pair within `rcut`, the interatomic distance is expanded on +an analytic radial basis (`basis_type`, with `n_radial` functions) and passed +through a radial network that produces one amplitude per channel. The ordered +pair of atom types modulates that amplitude with a learned scale and shift, so +the radial shape depends on which species face each other without instantiating +a separate network per pair. With `radial_modes` greater than zero, each ordered +pair additionally mixes several shared radial profiles with its own +coefficients, which lets pairs differ in shape rather than only in scale. + +Each amplitude is multiplied by a smooth cutoff envelope whose value and first +derivatives vanish at `rcut`, and by the real spherical harmonics of the +neighbor direction up to degree `lmax`. + +### Degree-wise moments and invariant read-out + +The per-atom state is the sum of these edge contributions, held separately for +each angular degree. Degree zero carries `channels` scalar values; higher +degrees carry progressively fewer channels, each with `2l + 1` angular +components. This tapering keeps the state small, which is what bounds both the +per-step cost and the memory per atom. + +The read-out contracts the moments into rotationally invariant scalars — norms +and cross-channel products within each degree, together with couplings across +degrees — and appends two measures of neighborhood density. Only these scalars +reach the fitting network: + +```math +\mathcal{D}_i = \mathrm{Invariants}\left(\{\mathbf{M}_i^{(l)}\}_{l=0}^{l_{\max}}\right). +``` + +Because the contraction is exactly rotationally invariant, the descriptor and +hence the energy are invariant under global rotation, and the forces obtained by +differentiation are equivariant. + +### Output calibration + +Descriptor statistics are used once, at initialization, to record a fixed +per-coordinate scale that puts the invariant outputs on a comparable footing +before they enter the fitting network. This is an initialization preconditioner, +not a running normalization: no sample-dependent statistic is evaluated during +training or inference, so the model remains a pure function of the atomic +positions and types. + +## Limitations + +- DPA4C is implemented for the PyTorch Exportable backend (`dp --pt-expt`). +- Export uses `.pt2` (AOTInductor); the TorchScript freeze path is not used. +- Model compression requires CUDA, `float32`, and a configuration inside the + compiled sets listed under [Model compression](#model-compression). +- The descriptor is one-hop local by construction. Interactions beyond `rcut` + are not represented, and unlike a message-passing model the effective range + cannot be extended by adding layers. diff --git a/doc/model/index.rst b/doc/model/index.rst index 8bd5aada64..d08b932059 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -12,6 +12,7 @@ Model dpa2 dpa3 dpa4 + dpa4c train-hybrid sel train-energy diff --git a/examples/water/dpa4c/README.md b/examples/water/dpa4c/README.md new file mode 100644 index 0000000000..5eff2964a7 --- /dev/null +++ b/examples/water/dpa4c/README.md @@ -0,0 +1,53 @@ +# Input for the DPA4C model + +This directory stores a configuration file for training DPA4C, the compact and +compressible degree-wise member of the DPA4 family. It runs on the pt_expt +backend: + +```bash +dp --pt-expt train input.json +``` + +DPA4C is built for extreme-speed molecular dynamics, so its arguments are best +read as a budget split between two quantities: inference throughput and the +largest system that fits in memory. + +## Descriptor arguments + +`channels` and `lmax` set every derived width. They are the only arguments that +grow the persistent equivariant node state, the per-atom tensor the descriptor +carries from the neighbor reduction into the readout, so raising either one +lowers both throughput and capacity. + +`radial_modes` lets each ordered atom-type pair combine several shared radial +shapes instead of rescaling a single one. It leaves the node state untouched +and spends per-edge work instead, which makes it the lever to reach for when +memory rather than throughput is the binding constraint. The compressed CUDA +path is compiled for the values `0`, `2`, `4` and `8`; a model trained with any +other value runs on the portable path but cannot be compressed. + +`basis_type` and `n_radial` select the analytic radial basis feeding the radial +network. + +`precision` must be `float32` for the compressed CUDA path. `use_amp` is an +execution policy rather than model state: it is read from this file for +training, while evaluation and inference read the `DP_AMP_INFER` environment +variable. The two are independent, so a model trained in full precision can +still be evaluated under mixed precision. + +## Choosing the fitting width + +The fitting network is sized against the descriptor, because the invariant +output grows with `channels`. The released grades pair `channels` 8, 32, 64 and +128 with hidden widths 96, 192, 256 and 384, at depth three as used here. This +file is the `Neo` grade: `channels: 32` with `radial_modes: 4` and a hidden +width of 192. + +Unlike `radial_modes`, fitting width is not a free trade against memory. It +does not enlarge the persistent node state, but it does add per-atom +activations and the derivatives saved for the force backward pass, so widening +or deepening it costs throughput and capacity together. + +Widen the fitting only when validation error is limited by fitting capacity +rather than by the descriptor; when either throughput or system size is +binding, spend the budget on the descriptor instead. diff --git a/examples/water/dpa4c/input.json b/examples/water/dpa4c/input.json new file mode 100644 index 0000000000..3b5e5e80df --- /dev/null +++ b/examples/water/dpa4c/input.json @@ -0,0 +1,78 @@ +{ + "_comment": "DPA4C energy-training example for the water dataset.", + "model": { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "_comment": "The Neo grade: channels and lmax fix every derived width, while radial_modes spends per-edge work without widening the per-atom state.", + "rcut": 6.0, + "channels": 32, + "lmax": 2, + "basis_type": "bessel", + "n_radial": 16, + "radial_modes": 4, + "_comment_precision": "float32 is required by the compressed CUDA path.", + "precision": "float32", + "use_amp": false, + "seed": 42 + }, + "fitting_net": { + "_comment": "The fitting width paired with channels 32 by the Neo grade.", + "neuron": [192, 192, 192], + "resnet_dt": false, + "activation_function": "silu", + "precision": "float32", + "seed": 42 + } + }, + "learning_rate": { + "type": "wsd", + "start_lr": 1e-3, + "stop_lr": 1e-6, + "warmup_ratio": 0.003, + "warmup_start_factor": 0.2, + "decay_phase_ratio": 0.2, + "decay_type": "inverse_linear" + }, + "loss": { + "type": "ener", + "loss_func": "mae", + "f_use_norm": true, + "start_pref_e": 20, + "limit_pref_e": 20, + "start_pref_f": 20, + "limit_pref_f": 20, + "start_pref_v": 5, + "limit_pref_v": 5 + }, + "optimizer": { + "type": "HybridMuon", + "weight_decay": 0.001 + }, + "training": { + "stat_file": "./dpa4c.hdf5", + "stat_file_mode": "update", + "training_data": { + "systems": ["../data/data_0", "../data/data_1", "../data/data_2"], + "batch_size": 1 + }, + "validation_data": { + "systems": ["../data/data_3"], + "batch_size": 1, + "numb_btch": 1 + }, + "numb_steps": 1000000, + "gradient_max_norm": 5.0, + "save_freq": 2000, + "max_ckpt_keep": 3, + "disp_file": "lcurve.out", + "disp_freq": 1000, + "disp_training": true, + "time_training": true, + "enable_compile": true, + "seed": 42 + }, + "validating": { + "save_best_dir": "ckpt_best" + } +} diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index b86ef773bb..97ecdb8dab 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -442,7 +442,7 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, * @param[in] nloc Number of owned local nodes. * @param[in] nall_nodes Total local-plus-halo node count. * @param[in] edge_storage Number of edge storage slots. - * @since API version 28 + * @since API version 29 */ extern void DP_DeepPotComputeCanonicalGraphGPU( DP_DeepPot* dp, @@ -450,11 +450,11 @@ extern void DP_DeepPotComputeCanonicalGraphGPU( double* d_force, double* d_atom_virial, const int64_t* d_atype, - const int64_t* d_source, + const uint32_t* d_source, const float* d_edge_vec, const int64_t* d_destination_row_ptr, const int64_t* d_source_row_ptr, - const int64_t* d_source_order, + const uint32_t* d_source_order, int nloc, int nall_nodes, int64_t edge_storage); diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 08cabbcf0c..fbff0c1823 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1351,11 +1351,11 @@ class DeepPot : public DeepBaseModel { double* d_force, double* d_atom_virial, const int64_t* d_atype, - const int64_t* d_source, + const uint32_t* d_source, const float* d_edge_vec, const int64_t* d_destination_row_ptr, const int64_t* d_source_row_ptr, - const int64_t* d_source_order, + const uint32_t* d_source_order, const int nloc, const int nall_nodes, const int64_t edge_storage) { diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index 4e91591ad3..e3f9ce63f0 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -1789,11 +1789,11 @@ void DP_DeepPotComputeCanonicalGraphGPU(DP_DeepPot* dp, double* d_force, double* d_atom_virial, const int64_t* d_atype, - const int64_t* d_source, + const uint32_t* d_source, const float* d_edge_vec, const int64_t* d_destination_row_ptr, const int64_t* d_source_row_ptr, - const int64_t* d_source_order, + const uint32_t* d_source_order, const int nloc, const int nall_nodes, const int64_t edge_storage) { diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 262c653713..8b2d7ae994 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -1,6 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #pragma once +#include #include #include "DeepBaseModel.h" @@ -357,11 +358,11 @@ class DeepPotBackend : public DeepBaseModelBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, - const std::int64_t* d_source, + const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, const std::int64_t* d_source_row_ptr, - const std::int64_t* d_source_order, + const std::uint32_t* d_source_order, const int nloc, const int nall_nodes, const std::int64_t edge_storage); @@ -829,11 +830,11 @@ class DeepPot : public DeepBaseModel { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, - const std::int64_t* d_source, + const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, const std::int64_t* d_source_row_ptr, - const std::int64_t* d_source_order, + const std::uint32_t* d_source_order, const int nloc, const int nall_nodes, const std::int64_t edge_storage); diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index a4ea314003..cde5b68cd9 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -348,11 +348,11 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, - const std::int64_t* d_source, + const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, const std::int64_t* d_source_row_ptr, - const std::int64_t* d_source_order, + const std::uint32_t* d_source_order, const int nloc, const int nall_nodes, const std::int64_t edge_storage) override; @@ -405,11 +405,11 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, - const std::int64_t* d_source, + const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, const std::int64_t* d_source_row_ptr, - const std::int64_t* d_source_order, + const std::uint32_t* d_source_order, const int nloc, const int nall_nodes, const std::int64_t edge_storage); diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 066c6ba991..388d48c3f6 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -399,21 +400,28 @@ inline CanonicalGraphTensorPack compactCanonicalGraph( graph.destination_row_ptr.select(0, graph.destination_row_ptr.size(0) - 1) .item(); const std::int64_t storage_count = std::max(edge_count, 2); + if (static_cast(storage_count) > + std::numeric_limits::max()) { + throw deepmd_exception( + "compact canonical graph exceeds the uint32 edge-index range"); + } auto source = torch::zeros({storage_count}, - graph.edge_index.options().dtype(torch::kInt64)); + graph.edge_index.options().dtype(torch::kUInt32)); auto edge_vec = torch::zeros({storage_count, 3}, graph.edge_vec.options().dtype(torch::kFloat32)); - auto source_order = torch::arange( - storage_count, graph.edge_index.options().dtype(torch::kInt64)); + auto source_order = + torch::arange(storage_count, + graph.edge_index.options().dtype(torch::kInt64)) + .to(torch::kUInt32); if (edge_count > 0) { source.slice(0, 0, edge_count) .copy_(graph.edge_index.select(0, 0) .slice(0, 0, edge_count) - .to(torch::kInt64)); + .to(torch::kUInt32)); edge_vec.slice(0, 0, edge_count) .copy_(graph.edge_vec.slice(0, 0, edge_count).to(torch::kFloat32)); source_order.slice(0, 0, edge_count) - .copy_(graph.source_order.slice(0, 0, edge_count).to(torch::kInt64)); + .copy_(graph.source_order.slice(0, 0, edge_count).to(torch::kUInt32)); } return {graph.atype, graph.n_node, diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 4888e0b387..dd9b57e2e3 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -661,11 +661,11 @@ void DeepPotBackend::compute_canonical_graph_gpu( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, - const std::int64_t* d_source, + const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, const std::int64_t* d_source_row_ptr, - const std::int64_t* d_source_order, + const std::uint32_t* d_source_order, const int nloc, const int nall_nodes, const std::int64_t edge_storage) { @@ -750,11 +750,11 @@ void DeepPot::compute_canonical_graph_gpu( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, - const std::int64_t* d_source, + const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, const std::int64_t* d_source_row_ptr, - const std::int64_t* d_source_order, + const std::uint32_t* d_source_order, const int nloc, const int nall_nodes, const std::int64_t edge_storage) { diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index b2ad13336c..4a85c18fe3 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -238,7 +239,8 @@ void DeepPotPTExpt::init(const std::string& model, metadata["lower_input_kind"].as_string(); lower_input_is_edge_ = lower_input_kind == "edge_vec"; lower_input_is_graph_ = lower_input_kind == "graph"; - lower_input_is_canonical_ = lower_input_kind == "dpa1_canonical"; + lower_input_is_canonical_ = lower_input_kind == "dpa1_canonical" || + lower_input_kind == "dpa4c_canonical"; } else { lower_input_is_edge_ = false; lower_input_is_graph_ = false; @@ -266,6 +268,12 @@ void DeepPotPTExpt::init(const std::string& model, throw deepmd::deepmd_exception( "compact canonical graph artifacts require float32 edge vectors."); } + if (!metadata.obj_val.count("canonical_index_dtype") || + metadata["canonical_index_dtype"].as_string() != "uint32") { + throw deepmd::deepmd_exception( + "compact canonical graph artifacts require uint32 topology; " + "re-freeze the model with the current DeePMD-kit version."); + } } type_map.clear(); @@ -2472,11 +2480,11 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, - const std::int64_t* d_source, + const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, const std::int64_t* d_source_row_ptr, - const std::int64_t* d_source_order, + const std::uint32_t* d_source_order, const int nloc, const int nall_nodes, const std::int64_t edge_storage) { @@ -2488,7 +2496,9 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( throw deepmd::deepmd_exception( "compute_canonical_graph_gpu requires a CUDA device."); } - if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2) { + if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2 || + static_cast(edge_storage) > + std::numeric_limits::max()) { throw deepmd::deepmd_exception( "invalid compact canonical graph dimensions."); } @@ -2500,10 +2510,12 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( torch::TensorOptions().dtype(torch::kFloat32).device(device); const auto opt_i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); + const auto opt_u32 = + torch::TensorOptions().dtype(torch::kUInt32).device(device); auto atype = torch::from_blob(const_cast(d_atype), {nall_nodes}, opt_i64); - auto source = torch::from_blob(const_cast(d_source), - {edge_storage}, opt_i64); + auto source = torch::from_blob(const_cast(d_source), + {edge_storage}, opt_u32); auto edge_vec = torch::from_blob(const_cast(d_edge_vec), {edge_storage, 3}, opt_f32); auto destination_row_ptr = @@ -2512,7 +2524,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( auto source_row_ptr = torch::from_blob( const_cast(d_source_row_ptr), {nall_nodes + 1}, opt_i64); auto source_order = torch::from_blob( - const_cast(d_source_order), {edge_storage}, opt_i64); + const_cast(d_source_order), {edge_storage}, opt_u32); auto n_node = torch::full({1}, nall_nodes, opt_i64); auto n_local = torch::full({1}, nloc, opt_i64); @@ -2602,11 +2614,11 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, - const std::int64_t* d_source, + const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, const std::int64_t* d_source_row_ptr, - const std::int64_t* d_source_order, + const std::uint32_t* d_source_order, const int nloc, const int nall_nodes, const std::int64_t edge_storage) { diff --git a/source/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc index ef83eae28b..f99ef18673 100644 --- a/source/api_cc/tests/test_neighbor_list_data.cc +++ b/source/api_cc/tests/test_neighbor_list_data.cc @@ -154,12 +154,12 @@ TEST(TestNeighborListData, CompactCanonicalGraphDropsMaskedGuards) { graph.source_row_ptr = torch::tensor({0, 1}, torch::kInt64); const auto compact = compactCanonicalGraph(graph); - EXPECT_EQ(compact.source.scalar_type(), torch::kInt64); + EXPECT_EQ(compact.source.scalar_type(), torch::kUInt32); EXPECT_EQ(compact.source.numel(), 2); EXPECT_EQ(compact.edge_vec.scalar_type(), torch::kFloat32); EXPECT_EQ(compact.edge_vec.size(0), 2); - EXPECT_TRUE( - torch::equal(compact.source_order, torch::tensor({0, 1}, torch::kInt64))); + EXPECT_TRUE(torch::equal(compact.source_order, + torch::tensor({0, 1}, torch::kUInt32))); EXPECT_EQ(compact.destination_row_ptr.select(0, 1).item(), 1); EXPECT_EQ(compact.source_row_ptr.select(0, 1).item(), 1); } diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index 9ecd68f055..67d604f863 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -27,6 +27,11 @@ using namespace LAMMPS_NS; +namespace { +// Lanes cooperating on one center's candidate list in the canonical fill. +constexpr int kNeighborLanes = 32; +} // namespace + template PairDeepMDKokkos::PairDeepMDKokkos(LAMMPS* lmp) : PairDeepMD(lmp), @@ -246,18 +251,15 @@ void PairDeepMDKokkos::prepare_model_nodes() { const int nall = atom->nlocal + atom->nghost; if (neighbor->ago == 0 || (int)k_loc2model.extent(0) < nall) { - if ((int)k_owner.extent(0) < nall) { - k_owner = DAT::tdual_int_1d("deepmd/kk:owner", nall); + if ((int)k_candidate_to_model.extent(0) < nall) { + k_candidate_to_model = + DAT::tdual_int_1d("deepmd/kk:candidate_to_model", nall); } if ((int)k_loc2model.extent(0) < nall) { k_loc2model = DAT::tdual_int_1d("deepmd/kk:loc2model", nall); k_model2loc = DAT::tdual_int_1d("deepmd/kk:model2loc", nall); } atomKK->sync(Host, TAG_MASK | TYPE_MASK); - auto h_owner = k_owner.view_host(); - for (int jj = 0; jj < nall; ++jj) { - h_owner(jj) = (jj < nlocal) ? jj : atom->map(atom->tag[jj]); - } auto h_loc2model = k_loc2model.view_host(); auto h_model2loc = k_model2loc.view_host(); const int* lmp_type = atom->type; @@ -282,9 +284,28 @@ void PairDeepMDKokkos::prepare_model_nodes() { } } nnode_model = m; - k_owner.template modify(); - k_owner.template sync(); - d_owner = k_owner.template view(); + + // Resolve each candidate atom to its model node once, on the host. In the + // folded representation a ghost contributes to the node of the local atom + // that owns it, so the resolution is a composition of the ownership map + // with the model map; the extended representation gives ghosts their own + // nodes and the composition degenerates to the model map. Collapsing it + // here leaves the device traversal, which visits every candidate of every + // center, with a single gather. + auto h_candidate_to_model = k_candidate_to_model.view_host(); + if (multi_rank) { + for (int j = 0; j < nall; ++j) { + h_candidate_to_model(j) = h_loc2model(j); + } + } else { + for (int j = 0; j < nall; ++j) { + const int owner = (j < nlocal) ? j : atom->map(atom->tag[j]); + h_candidate_to_model(j) = owner < 0 ? -1 : h_loc2model(owner); + } + } + k_candidate_to_model.template modify(); + k_candidate_to_model.template sync(); + d_candidate_to_model = k_candidate_to_model.template view(); k_loc2model.template modify(); k_loc2model.template sync(); d_loc2model = k_loc2model.template view(); @@ -340,9 +361,8 @@ int PairDeepMDKokkos::build_edges_device() { const double cut = cutoff; const double cutsq = cut * cut; - const bool multi = multi_rank; - auto owner = d_owner; auto loc2model = d_loc2model; + auto candidate_to_model = d_candidate_to_model; auto model2loc = d_model2loc; if ((int)d_edge_offset.extent(0) < nlocal + 1) { @@ -368,7 +388,7 @@ int PairDeepMDKokkos::build_edges_device() { for (int jj = 0; jj < jnum; ++jj) { int j = d_neighbors(i, jj); j &= NEIGHMASK; - if (loc2model(multi ? j : owner(j)) < 0) { + if (candidate_to_model(j) < 0) { continue; } const double dx = x(j, 0) - xi, dy = x(j, 1) - yi, dz = x(j, 2) - zi; @@ -451,7 +471,7 @@ int PairDeepMDKokkos::build_edges_device() { for (int jj = 0; jj < jnum; ++jj) { int j = d_neighbors(i, jj); j &= NEIGHMASK; - const int mj = loc2model(multi ? j : owner(j)); + const int mj = candidate_to_model(j); if (mj < 0) { continue; } @@ -505,9 +525,8 @@ std::int64_t PairDeepMDKokkos::build_canonical_edges_device( auto d_ilist = k_list->d_ilist; atomKK->sync(execution_space, X_MASK); auto x = atomKK->k_x.template view(); - auto owner = d_owner; auto loc2model = d_loc2model; - const bool multi = multi_rank; + auto candidate_to_model = d_candidate_to_model; const double cutsq = cutoff * cutoff; const double inv_dist = 1.0 / dist_unit_cvt_factor; const int node_count_int = nnode_model; @@ -516,15 +535,15 @@ std::int64_t PairDeepMDKokkos::build_canonical_edges_device( if (workspace.destination_row_ptr.extent(0) < node_count + 1) { workspace.destination_row_ptr = Kokkos::View( "deepmd/kk:canonical_destination_row_ptr", node_count + 1); - workspace.source_counts = Kokkos::View( + workspace.source_counts = Kokkos::View( "deepmd/kk:canonical_source_counts", node_count); workspace.source_row_ptr = Kokkos::View( "deepmd/kk:canonical_source_row_ptr", node_count + 1); - workspace.source_cursor = Kokkos::View( + workspace.source_cursor = Kokkos::View( "deepmd/kk:canonical_source_cursor", node_count); } Kokkos::deep_copy(workspace.destination_row_ptr, std::int64_t{0}); - Kokkos::deep_copy(workspace.source_counts, std::int64_t{0}); + Kokkos::deep_copy(workspace.source_counts, std::uint32_t{0}); if (node_count_int == 0) { return 0; } @@ -545,8 +564,8 @@ std::int64_t PairDeepMDKokkos::build_canonical_edges_device( const int jnum = d_numneigh(i); std::int64_t count = 0; for (int jj = 0; jj < jnum; ++jj) { - int j = d_neighbors(i, jj) & NEIGHMASK; - const int mj = loc2model(multi ? j : owner(j)); + const int j = d_neighbors(i, jj) & NEIGHMASK; + const int mj = candidate_to_model(j); if (mj < 0) { continue; } @@ -577,53 +596,112 @@ std::int64_t PairDeepMDKokkos::build_canonical_edges_device( Kokkos::deep_copy(edge_count, Kokkos::subview(workspace.destination_row_ptr, node_count_int)); const std::int64_t storage_count = std::max(edge_count, 2); + if (static_cast(storage_count) > + std::numeric_limits::max()) { + error->one(FLERR, + "Compact canonical graph exceeds the uint32 edge-index range"); + } const std::size_t required = static_cast(storage_count); if (workspace.edge_capacity < required) { - const std::size_t slack = required / 8 + 64; + // Thermal cutoff-count fluctuations are much smaller than the historical + // 12.5% geometric-growth reserve. A 2% reserve avoids repeated allocation + // while preventing unused edge storage from retaining several GiB at + // billion-edge scale. + const std::size_t slack = required / 50 + 64; if (required > std::numeric_limits::max() - slack) { - error->one(FLERR, "Compact DPA1 graph capacity overflows size_t"); + error->one(FLERR, "Compact canonical graph capacity overflows size_t"); } workspace.edge_capacity = required + slack; - workspace.source = Kokkos::View( + workspace.source = Kokkos::View( "deepmd/kk:canonical_source", workspace.edge_capacity); workspace.edge_vec = Kokkos::View( "deepmd/kk:canonical_edge_vec", workspace.edge_capacity * 3); - workspace.source_order = Kokkos::View( + workspace.source_order = Kokkos::View( "deepmd/kk:canonical_source_order", workspace.edge_capacity); } auto source = workspace.source; auto edge_vec = workspace.edge_vec; + // One warp per center. A thread-per-center fill writes each surviving edge + // at an offset private to its center, so the lanes of a warp scatter their + // twelve-byte edge vectors across thirty-two unrelated rows. Cooperating on + // one center instead sends consecutive survivors to consecutive slots, which + // coalesces the dominant store stream. Candidates are taken a warp at a + // time and each lane writes at its exclusive prefix within the warp, so the + // edge order is the candidate order the serial fill produces. + using team_policy = Kokkos::TeamPolicy; + using member_type = typename team_policy::member_type; + using lane_scratch = + Kokkos::View>; + using vector_scratch = + Kokkos::View>; + const int scratch_bytes = lane_scratch::shmem_size(kNeighborLanes) + + vector_scratch::shmem_size(3 * kNeighborLanes); Kokkos::parallel_for( - "deepmd/kk:canonical_fill", Kokkos::RangePolicy(0, inum), - KOKKOS_LAMBDA(const int ii) { - const int i = d_ilist(ii); + "deepmd/kk:canonical_fill", + team_policy(inum, kNeighborLanes) + .set_scratch_size(0, Kokkos::PerTeam(scratch_bytes)), + KOKKOS_LAMBDA(const member_type& team) { + const int i = d_ilist(team.league_rank()); const int mi = loc2model(i); if (mi < 0) { return; } + lane_scratch node(team.team_scratch(0), kNeighborLanes); + vector_scratch vec(team.team_scratch(0), 3 * kNeighborLanes); + const double xi = x(i, 0); const double yi = x(i, 1); const double zi = x(i, 2); const int jnum = d_numneigh(i); + const int lane = team.team_rank(); std::int64_t edge = destination_row_ptr(mi); - for (int jj = 0; jj < jnum; ++jj) { - int j = d_neighbors(i, jj) & NEIGHMASK; - const int mj = loc2model(multi ? j : owner(j)); - if (mj < 0) { - continue; - } - const double dx = x(j, 0) - xi; - const double dy = x(j, 1) - yi; - const double dz = x(j, 2) - zi; - if (dx * dx + dy * dy + dz * dz < cutsq) { - source(edge) = static_cast(mj); - edge_vec(3 * edge + 0) = static_cast(dx * inv_dist); - edge_vec(3 * edge + 1) = static_cast(dy * inv_dist); - edge_vec(3 * edge + 2) = static_cast(dz * inv_dist); - Kokkos::atomic_fetch_add(&source_counts(mj), std::int64_t{1}); - ++edge; + for (int base = 0; base < jnum; base += kNeighborLanes) { + const int jj = base + lane; + int mj = -1; + if (jj < jnum) { + const int j = d_neighbors(i, jj) & NEIGHMASK; + mj = candidate_to_model(j); + if (mj >= 0) { + const double dx = x(j, 0) - xi; + const double dy = x(j, 1) - yi; + const double dz = x(j, 2) - zi; + if (dx * dx + dy * dy + dz * dz < cutsq) { + vec(3 * lane + 0) = static_cast(dx * inv_dist); + vec(3 * lane + 1) = static_cast(dy * inv_dist); + vec(3 * lane + 2) = static_cast(dz * inv_dist); + } else { + mj = -1; + } + } } + node(lane) = mj; + team.team_barrier(); + + // Compact the survivors of this warp of candidates: consecutive + // survivors take consecutive slots, so the stores of a warp fall in + // one contiguous span of the edge arrays. + std::int64_t kept = 0; + Kokkos::parallel_scan( + Kokkos::TeamThreadRange(team, kNeighborLanes), + [&](const int slot, std::int64_t& offset, const bool final) { + const int target = node(slot); + if (final && target >= 0) { + const std::int64_t position = edge + offset; + source(position) = static_cast(target); + edge_vec(3 * position + 0) = vec(3 * slot + 0); + edge_vec(3 * position + 1) = vec(3 * slot + 1); + edge_vec(3 * position + 2) = vec(3 * slot + 2); + Kokkos::atomic_fetch_add(&source_counts(target), + std::uint32_t{1}); + } + offset += target >= 0 ? 1 : 0; + }, + kept); + edge += kept; + team.team_barrier(); } }); @@ -632,7 +710,8 @@ std::int64_t PairDeepMDKokkos::build_canonical_edges_device( "deepmd/kk:canonical_source_scan", Kokkos::RangePolicy(0, node_count_int), KOKKOS_LAMBDA(const int node, std::int64_t& update, const bool final) { - const std::int64_t count = source_counts(node); + const std::int64_t count = + static_cast(source_counts(node)); if (final) { source_row_ptr(node) = update; } @@ -641,21 +720,22 @@ std::int64_t PairDeepMDKokkos::build_canonical_edges_device( source_row_ptr(node_count_int) = update; } }); - Kokkos::deep_copy( - workspace.source_cursor, - Kokkos::subview(workspace.source_row_ptr, - std::make_pair(std::int64_t{0}, static_cast( - node_count_int)))); auto source_cursor = workspace.source_cursor; + Kokkos::parallel_for( + "deepmd/kk:canonical_source_cursor", + Kokkos::RangePolicy(0, node_count_int), + KOKKOS_LAMBDA(const int node) { + source_cursor(node) = static_cast(source_row_ptr(node)); + }); auto source_order = workspace.source_order; Kokkos::parallel_for( "deepmd/kk:canonical_source_scatter", Kokkos::RangePolicy>( 0, edge_count), KOKKOS_LAMBDA(const std::int64_t edge) { - const auto position = - Kokkos::atomic_fetch_add(&source_cursor(source(edge)), 1LL); - source_order(position) = edge; + const auto position = Kokkos::atomic_fetch_add( + &source_cursor(source(edge)), std::uint32_t{1}); + source_order(position) = static_cast(edge); }); if (storage_count > edge_count) { Kokkos::parallel_for( @@ -663,11 +743,11 @@ std::int64_t PairDeepMDKokkos::build_canonical_edges_device( Kokkos::RangePolicy>( edge_count, storage_count), KOKKOS_LAMBDA(const std::int64_t edge) { - source(edge) = 0; + source(edge) = std::uint32_t{0}; edge_vec(3 * edge + 0) = 0.0f; edge_vec(3 * edge + 1) = 0.0f; edge_vec(3 * edge + 2) = 0.0f; - source_order(edge) = edge; + source_order(edge) = static_cast(edge); }); } return edge_count; diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index d09e9a7abc..93cfe4c358 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -29,13 +29,15 @@ namespace LAMMPS_NS { template struct CompactCanonicalGraphWorkspace { - Kokkos::View source; + Kokkos::View source; Kokkos::View edge_vec; Kokkos::View destination_row_ptr; - Kokkos::View source_counts; + // Per-node counts are bounded by the LAMMPS neighbor limit. CSR offsets + // remain int64 so the compact graph retains its full global edge range. + Kokkos::View source_counts; Kokkos::View source_row_ptr; - Kokkos::View source_cursor; - Kokkos::View source_order; + Kokkos::View source_cursor; + Kokkos::View source_order; std::size_t edge_capacity = 0; }; @@ -97,10 +99,6 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { d_model_type; // (nnode_model) type per model node Kokkos::View d_model_type_i64; // compact canonical artifact type per model node - // Ghost -> local owner fold, rebuilt on the host at each neighbor rebuild. - DAT::tdual_int_1d k_owner; - typename AT::t_int_1d d_owner; - // Virtual-atom (NULL type) compaction, rebuilt with the neighbor list: the // model sees only the local atoms with a real model type, so ``model2loc`` // lists those local indices and ``loc2model`` inverts it (-1 for virtual). @@ -110,6 +108,11 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { int nloc_model; // real local model nodes; the energy is summed over these int nnode_model; // total model nodes (== nloc_model folded; + ghost // extended) + // (nall) candidate atom -> model node index, or -1. Folding a ghost onto its + // owner and mapping that atom to a node is resolved once per neighbor + // rebuild so that the graph traversal needs a single gather per candidate. + DAT::tdual_int_1d k_candidate_to_model; + typename AT::t_int_1d d_candidate_to_model; DAT::tdual_int_1d k_loc2model; // (nall) atom -> model node index, or -1 DAT::tdual_int_1d k_model2loc; // (nall) model node index -> atom index typename AT::t_int_1d d_loc2model; diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 5e1a46ec91..adb3bd4aad 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -35,6 +35,12 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) dpa1_graph_descriptor.cu dpa1_graph_compress.cu ${DPA1_GRAPH_COMPRESS_KERNEL_SRC} + dpa4c_graph_compress.cu + dpa4c_graph_compress_c8.cu + dpa4c_graph_compress_c16.cu + dpa4c_graph_compress_c32.cu + dpa4c_graph_compress_c64.cu + dpa4c_graph_compress_c128.cu graph_fitting.cu edge_force_virial.cu dpa1_graph_energy_force.cu) @@ -68,11 +74,14 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) target_compile_definitions(deepmd_op_pt PRIVATE DEEPMD_ENABLE_DPA1_HIGH_LMAX=1) endif() - # The compressed DPA1 kernels are instantiated one translation unit per - # channel width so their topology and angular-degree variants compile in - # parallel. - set_source_files_properties(${DPA1_GRAPH_COMPRESS_KERNEL_SRC} - PROPERTIES COMPILE_OPTIONS "--use_fast_math") + # The compressed DPA1 and DPA4C kernels are instantiated one translation + # unit per channel width so their angular-degree and topology specializations + # compile in parallel. + set_source_files_properties( + ${DPA1_GRAPH_COMPRESS_KERNEL_SRC} dpa4c_graph_compress_c8.cu + dpa4c_graph_compress_c16.cu dpa4c_graph_compress_c32.cu + dpa4c_graph_compress_c64.cu dpa4c_graph_compress_c128.cu + PROPERTIES COMPILE_OPTIONS "--use_fast_math") endif() if(${OP_CXX_ABI_PT} EQUAL ${OP_CXX_ABI}) target_link_libraries(deepmd_op_pt PRIVATE ${LIB_DEEPMD}) diff --git a/source/op/pt/dpa1_graph_compress.cu b/source/op/pt/dpa1_graph_compress.cu index d41c833092..53fe4800a0 100644 --- a/source/op/pt/dpa1_graph_compress.cu +++ b/source/op/pt/dpa1_graph_compress.cu @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -46,6 +47,29 @@ void check_launch(const char* operation, const cudaError_t error) { TORCH_CHECK(error == cudaSuccess, operation, ": ", cudaGetErrorString(error)); } +IndexKind index_kind_from_tensor(const torch::Tensor& tensor) { + if (tensor.scalar_type() == torch::kInt32) { + return IndexKind::kInt32; + } + if (tensor.scalar_type() == torch::kUInt32) { + return IndexKind::kUInt32; + } + return IndexKind::kInt64; +} + +const void* index_data_ptr(const torch::Tensor& tensor, IndexKind kind) { + switch (kind) { + case IndexKind::kInt32: + return static_cast(tensor.data_ptr()); + case IndexKind::kUInt32: + return static_cast(tensor.data_ptr()); + case IndexKind::kInt64: + return static_cast(tensor.data_ptr()); + } + TORCH_CHECK(false, "dpa1_graph_compress: unsupported index kind"); + return nullptr; +} + Arguments make_common_arguments(long node_count, int basis_dim, int type_count, @@ -91,9 +115,7 @@ Arguments make_common_arguments(long node_count, arguments.smooth = smooth; arguments.canonical = canonical; arguments.masked = edge_mask.numel() != 0; - arguments.index_kind = edge_index.scalar_type() == torch::kInt32 - ? IndexKind::kInt32 - : IndexKind::kInt64; + arguments.index_kind = index_kind_from_tensor(edge_index); arguments.rcut = rcut; arguments.rcut_smooth = rcut_smooth; arguments.protection = protection; @@ -104,18 +126,12 @@ Arguments make_common_arguments(long node_count, arguments.stride0 = stride0; arguments.stride1 = stride1; arguments.edge_vec = edge_vec.data_ptr(); - arguments.edge_index = - arguments.index_kind == IndexKind::kInt32 - ? static_cast(edge_index.data_ptr()) - : static_cast(edge_index.data_ptr()); + arguments.edge_index = index_data_ptr(edge_index, arguments.index_kind); arguments.edge_mask = arguments.masked ? edge_mask.data_ptr() : nullptr; arguments.destination_order = destination_order.numel() == 0 ? nullptr - : (arguments.index_kind == IndexKind::kInt32 - ? static_cast(destination_order.data_ptr()) - : static_cast( - destination_order.data_ptr())); + : index_data_ptr(destination_order, arguments.index_kind); arguments.destination_row_ptr = destination_row_ptr.data_ptr(); arguments.atype = atype.data_ptr(); arguments.average = average.data_ptr(); @@ -153,8 +169,10 @@ void validate_inputs(const torch::Tensor& edge_vec, table.is_contiguous() && gate_table.is_contiguous(), "dpa1_graph_compress: inputs must be contiguous"); TORCH_CHECK(edge_index.scalar_type() == torch::kInt32 || + edge_index.scalar_type() == torch::kUInt32 || edge_index.scalar_type() == torch::kInt64, - "dpa1_graph_compress: edge_index must be int32 or int64"); + "dpa1_graph_compress: edge_index must be int32, uint32, or " + "int64"); TORCH_CHECK(destination_order.scalar_type() == edge_index.scalar_type(), "dpa1_graph_compress: destination_order must match the " "edge_index dtype"); diff --git a/source/op/pt/dpa1_graph_compress_kernel.cuh b/source/op/pt/dpa1_graph_compress_kernel.cuh index 91362858d8..50c38c8ba3 100644 --- a/source/op/pt/dpa1_graph_compress_kernel.cuh +++ b/source/op/pt/dpa1_graph_compress_kernel.cuh @@ -28,6 +28,7 @@ #include #include +#include #include "dpa1_graph_compress_launch.h" @@ -1141,6 +1142,9 @@ cudaError_t dispatch_forward_index(const Arguments& arguments, if (arguments.index_kind == IndexKind::kInt32) { return dispatch_forward_basis(arguments, stream); } + if (arguments.index_kind == IndexKind::kUInt32) { + return dispatch_forward_basis(arguments, stream); + } return dispatch_forward_basis(arguments, stream); } @@ -1150,6 +1154,9 @@ cudaError_t dispatch_backward_index(const Arguments& arguments, if (arguments.index_kind == IndexKind::kInt32) { return dispatch_backward_basis(arguments, stream); } + if (arguments.index_kind == IndexKind::kUInt32) { + return dispatch_backward_basis(arguments, stream); + } return dispatch_backward_basis(arguments, stream); } diff --git a/source/op/pt/dpa1_graph_compress_launch.h b/source/op/pt/dpa1_graph_compress_launch.h index 40e4e12e92..72a8e7960e 100644 --- a/source/op/pt/dpa1_graph_compress_launch.h +++ b/source/op/pt/dpa1_graph_compress_launch.h @@ -18,6 +18,7 @@ namespace deepmd_dpa1_compress { enum class IndexKind : int { kInt32 = 0, kInt64 = 1, + kUInt32 = 2, }; struct Arguments { diff --git a/source/op/pt/dpa1_graph_energy_force.cu b/source/op/pt/dpa1_graph_energy_force.cu index 04942e1493..212e4a7550 100644 --- a/source/op/pt/dpa1_graph_energy_force.cu +++ b/source/op/pt/dpa1_graph_energy_force.cu @@ -68,7 +68,6 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, int64_t basis_dim, std::vector fit_ws, std::vector fit_bs, - std::vector fit_idts, std::vector fit_resnets, torch::Tensor w_head, torch::Tensor b_head, @@ -102,7 +101,7 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, const torch::Tensor& g_saved = std::get<6>(desc); // === Step 2. Fitting forward: descriptor -> per-atom energy. === - auto fit = graph_fitting(grrg, atype, fit_ws, fit_bs, fit_idts, fit_resnets, + auto fit = graph_fitting(grrg, atype, fit_ws, fit_bs, fit_resnets, w_head, b_head, bias_atom_e, fit_act); const torch::Tensor& atom_energy_raw = std::get<0>(fit); // (N, 1) fp64 const torch::Tensor& fit_saved = std::get<1>(fit); @@ -119,8 +118,8 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, // === Step 4. Force = grad of the reduced energy; dE_redu/d(atom_e) == 1. === std::get<0>(desc) = torch::Tensor(); - auto d_grrg = graph_fitting_backward(energy_seed, fit_saved, fit_ws, - fit_resnets, w_head); + auto d_grrg = graph_fitting_backward(energy_seed, fit_saved, fit_ws, fit_bs, + fit_resnets, w_head, fit_act); std::get<1>(fit) = torch::Tensor(); auto g_e = dpa1_graph_descriptor_backward( d_grrg, std::nullopt, gr, edge_order, pair_table, pre2_saved, g_saved, @@ -151,9 +150,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, m) { "b2, Tensor idt2, Tensor w3, Tensor b3, Tensor idt3, Tensor gate_table, " "int act, int type_one_side, int concat_tebd, int smooth, int axis, int " "resnet2, int resnet3, float rcut, float rcut_smth, float protection, " - "float nnei, int basis_dim, Tensor[] fit_ws, Tensor[] fit_bs, Tensor[] " - "fit_idts, int[] " - "fit_resnets, Tensor w_head, Tensor b_head, Tensor bias_atom_e, int " + "float nnei, int basis_dim, Tensor[] fit_ws, Tensor[] fit_bs, " + "int[] fit_resnets, Tensor w_head, Tensor b_head, Tensor bias_atom_e, int " "fit_act, SymInt node_capacity, bool do_atomic_virial) -> (Tensor, " "Tensor, Tensor, Tensor, Tensor)"); m.impl("dpa1_graph_energy_force", torch::kCUDA, &dpa1_graph_energy_force); diff --git a/source/op/pt/dpa4c_graph_compress.cu b/source/op/pt/dpa4c_graph_compress.cu new file mode 100644 index 0000000000..74db36f8de --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress.cu @@ -0,0 +1,794 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Torch bindings of the compressed degree-wise DPA4C descriptor. +// +// This translation unit owns argument validation, the scalar-width dispatch, +// and the operator registration. The kernels themselves are instantiated in +// one translation unit per scalar width. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "dpa4c_graph_compress_launch.h" +#include "graph_ops.h" + +namespace { + +using deepmd_dpa4c::Arguments; +using deepmd_dpa4c::IndexKind; +using deepmd_dpa4c::Profile; + +#define DPA4C_CHECK_LAUNCH(name) \ + do { \ + const cudaError_t error = cudaGetLastError(); \ + TORCH_CHECK(error == cudaSuccess, name, ": ", cudaGetErrorString(error)); \ + } while (0) + +/// Every width the host must reproduce to validate one operator invocation. +struct Dimensions { + int moment_width; + int output_width; + int degree_one; + int coupling_records; +}; + +template +Dimensions dimensions_of() { + using P = Profile; + return {P::MomentWidth, P::OutputWidth, P::C1, + deepmd_dpa4c::coupling_record_count(Lmax)}; +} + +template +Dimensions dimensions_for(int lmax) { + switch (lmax) { + case 2: + return dimensions_of(); + case 3: + return dimensions_of(); + case 4: + return dimensions_of(); + default: + TORCH_CHECK(false, "dpa4c_graph_compress: unsupported lmax ", lmax); + } +} + +Dimensions profile_dimensions(int channels, int lmax) { +#define DPA4C_DIMENSIONS(width) \ + if (channels == width) { \ + return dimensions_for(lmax); \ + } + DPA4C_FOR_EACH_CHANNEL(DPA4C_DIMENSIONS) +#undef DPA4C_DIMENSIONS + TORCH_CHECK(false, "dpa4c_graph_compress: unsupported channels ", channels); +} + +void dispatch(int channels, + bool backward, + const Arguments& arguments, + cudaStream_t stream) { +#define DPA4C_DISPATCH(width) \ + if (channels == width) { \ + if (backward) { \ + deepmd_dpa4c::launch_backward_c##width(arguments, stream); \ + } else { \ + deepmd_dpa4c::launch_forward_c##width(arguments, stream); \ + } \ + DPA4C_CHECK_LAUNCH("dpa4c_graph_compress"); \ + return; \ + } + DPA4C_FOR_EACH_CHANNEL(DPA4C_DISPATCH) +#undef DPA4C_DISPATCH + TORCH_CHECK(false, "dpa4c_graph_compress: unsupported channels ", channels); +} + +/// Immutable tensors and scalars shared by both directions of the operator. +struct Payload { + torch::Tensor edge_index; + torch::Tensor edge_mask; + torch::Tensor destination_order; + torch::Tensor destination_row_ptr; + torch::Tensor atype; + torch::Tensor table; + torch::Tensor pair_film; + torch::Tensor pair_mixing; + torch::Tensor type_embedding; + torch::Tensor readout_matrices; + torch::Tensor coupling_meta; + torch::Tensor coupling_entry; + torch::Tensor coupling_value; + torch::Tensor output_mean; + torch::Tensor output_inv_std; + bool canonical; + int64_t lmax; + double table_stride; + double table_max; + double rcut; + double eps; + double degree_floor; + int64_t node_begin = 0; +}; + +IndexKind index_kind_of(const torch::Tensor& edge_index) { + switch (edge_index.scalar_type()) { + case torch::kInt32: + case torch::kUInt32: + return IndexKind::Bits32; + case torch::kInt64: + return IndexKind::Bits64; + default: + TORCH_CHECK(false, + "dpa4c_graph_compress: edge indices must be int32, uint32, " + "or int64"); + } +} + +Arguments build_arguments(const Payload& payload, + const torch::Tensor& edge_vec, + int channels, + const Dimensions& widths) { + // A tiled caller passes a slice of the destination row pointer, so the run + // length follows the slice while the type table stays system-wide. + const long node_count = payload.destination_row_ptr.numel() - 1; + const int type_count = static_cast(payload.type_embedding.size(0)); + const int radial_modes = payload.pair_mixing.numel() == 0 + ? 0 + : static_cast(payload.pair_mixing.size(2)); + + const torch::Device device = edge_vec.device(); + for (const torch::Tensor* tensor : + {&edge_vec, &payload.edge_index, &payload.edge_mask, + &payload.destination_order, &payload.destination_row_ptr, + &payload.atype, &payload.table, &payload.pair_film, + &payload.pair_mixing, &payload.type_embedding, + &payload.readout_matrices, &payload.coupling_meta, + &payload.coupling_entry, &payload.coupling_value, &payload.output_mean, + &payload.output_inv_std}) { + TORCH_CHECK(tensor->is_cuda() && tensor->device() == device, + "dpa4c_graph_compress: every tensor input must be a CUDA " + "tensor on the device of edge_vec"); + TORCH_CHECK(tensor->is_contiguous(), + "dpa4c_graph_compress: all tensor inputs must be contiguous"); + } + TORCH_CHECK(payload.destination_order.scalar_type() == + payload.edge_index.scalar_type(), + "dpa4c_graph_compress: destination_order dtype must match " + "edge_index"); + TORCH_CHECK(payload.edge_mask.scalar_type() == torch::kBool, + "dpa4c_graph_compress: edge_mask must be bool"); + TORCH_CHECK( + payload.destination_row_ptr.scalar_type() == torch::kInt64 && + payload.atype.scalar_type() == torch::kInt64, + "dpa4c_graph_compress: row pointers and atom types must be int64"); + TORCH_CHECK(payload.coupling_meta.scalar_type() == torch::kInt32 && + payload.coupling_entry.scalar_type() == torch::kInt32, + "dpa4c_graph_compress: the coupling layout must be int32"); + for (const torch::Tensor* tensor : + {&payload.table, &payload.pair_film, &payload.pair_mixing, + &payload.type_embedding, &payload.readout_matrices, + &payload.coupling_value, &payload.output_mean, + &payload.output_inv_std}) { + TORCH_CHECK(tensor->scalar_type() == torch::kFloat32, + "dpa4c_graph_compress: tables and weights must be fp32"); + } + TORCH_CHECK(payload.lmax >= 2 && payload.lmax <= 4, + "dpa4c_graph_compress: lmax must be 2, 3, or 4"); + // The shared mode cache is sized from a compile-time maximum and the split + // spline row assumes an even table width, so the rank set is closed. + TORCH_CHECK(radial_modes == 0 || radial_modes == 2 || radial_modes == 4 || + radial_modes == 8, + "dpa4c_graph_compress: radial_modes must be 0, 2, 4, or 8, got ", + radial_modes); + TORCH_CHECK(edge_vec.dim() == 2 && edge_vec.size(1) == 3, + "dpa4c_graph_compress: edge_vec must have shape (E, 3)"); + TORCH_CHECK(payload.table.dim() == 2 && payload.table.size(0) > 0 && + payload.table.size(1) == 6 * (channels + radial_modes), + "dpa4c_graph_compress: invalid radial table shape"); + TORCH_CHECK(payload.type_embedding.dim() == 2 && type_count > 1, + "dpa4c_graph_compress: invalid type embedding shape"); + TORCH_CHECK( + payload.pair_film.sizes() == + torch::IntArrayRef( + {static_cast(type_count) * type_count, channels, 2}), + "dpa4c_graph_compress: invalid PairFiLM cache shape"); + TORCH_CHECK( + radial_modes == 0 || + payload.pair_mixing.sizes() == + torch::IntArrayRef({static_cast(type_count) * type_count, + channels, radial_modes}), + "dpa4c_graph_compress: invalid mode-mixing cache shape"); + // The kernel addresses the packed projections with a fixed degree-one + // stride, so the padded block must have exactly that extent. + TORCH_CHECK(payload.readout_matrices.sizes() == + torch::IntArrayRef({8, widths.degree_one, widths.degree_one}), + "dpa4c_graph_compress: invalid readout matrix shape"); + TORCH_CHECK(payload.coupling_meta.dim() == 2 && + payload.coupling_meta.size(1) == 8 && + payload.coupling_meta.size(0) == widths.coupling_records, + "dpa4c_graph_compress: the coupling layout must describe ", + widths.coupling_records, " degree triples"); + TORCH_CHECK(payload.coupling_entry.numel() == payload.coupling_value.numel(), + "dpa4c_graph_compress: coupling coordinates and values must " + "have equal length"); + TORCH_CHECK(payload.output_mean.numel() == widths.output_width && + payload.output_inv_std.numel() == widths.output_width, + "dpa4c_graph_compress: invalid output calibration shape"); + // A run covers ``node_count`` destination rows starting at ``node_begin``; + // the type table is system-wide because neighbor lookups index it with + // absolute source indices. Whole-system entry points additionally require + // the two to describe the same node axis. + TORCH_CHECK(payload.atype.size(0) >= payload.node_begin + node_count, + "dpa4c_graph_compress: destination_row_ptr must have N + 1 " + "entries"); + + Arguments arguments; + arguments.node_count = node_count; + arguments.edge_count = edge_vec.size(0); + arguments.lmax = static_cast(payload.lmax); + arguments.interval_count = static_cast(payload.table.size(0)); + arguments.type_count = type_count; + arguments.table_width = channels + radial_modes; + arguments.radial_modes = radial_modes; + arguments.coupling_count = static_cast(payload.coupling_meta.size(0)); + arguments.table_stride = static_cast(payload.table_stride); + arguments.table_max = static_cast(payload.table_max); + arguments.rcut = static_cast(payload.rcut); + arguments.eps = static_cast(payload.eps); + arguments.degree_floor = static_cast(payload.degree_floor); + arguments.canonical = payload.canonical; + arguments.index_kind = index_kind_of(payload.edge_index); + arguments.edge_index = payload.edge_index.data_ptr(); + arguments.destination_order = payload.destination_order.numel() != 0 + ? payload.destination_order.data_ptr() + : nullptr; + arguments.edge_vec = edge_vec.data_ptr(); + arguments.edge_mask = payload.edge_mask.numel() != 0 + ? payload.edge_mask.data_ptr() + : nullptr; + arguments.destination_row_ptr = payload.destination_row_ptr.data_ptr(); + arguments.node_begin = payload.node_begin; + arguments.atype = payload.atype.data_ptr(); + arguments.table = payload.table.data_ptr(); + arguments.pair_film = payload.pair_film.data_ptr(); + arguments.pair_mixing = + radial_modes != 0 ? payload.pair_mixing.data_ptr() : nullptr; + arguments.type_embedding = payload.type_embedding.data_ptr(); + arguments.readout_matrices = payload.readout_matrices.data_ptr(); + arguments.coupling_meta = payload.coupling_meta.data_ptr(); + arguments.coupling_entry = payload.coupling_entry.data_ptr(); + arguments.coupling_value = payload.coupling_value.data_ptr(); + arguments.output_mean = payload.output_mean.data_ptr(); + arguments.output_inv_std = payload.output_inv_std.data_ptr(); + return arguments; +} + +} // namespace + +std::tuple dpa4c_graph_compress( + torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor destination_order, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + bool canonical, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor) { + const Payload payload{edge_index, edge_mask, destination_order, + destination_row_ptr, atype, table, + pair_film, pair_mixing, type_embedding, + readout_matrices, coupling_meta, coupling_entry, + coupling_value, output_mean, output_inv_std, + canonical, lmax, table_stride, + table_max, rcut, eps, + degree_floor}; + const long node_count = destination_row_ptr.numel() - 1; + // The destination row pointer defines the node axis; the type table must + // describe exactly that axis, or the two disagree on how many nodes exist. + TORCH_CHECK(atype.size(0) == destination_row_ptr.numel() - 1, + "dpa4c_graph_compress: atype and destination_row_ptr describe " + "different node counts"); + const int channels = static_cast(type_embedding.size(1)); + const Dimensions widths = + profile_dimensions(channels, static_cast(lmax)); + TORCH_CHECK(edge_vec.is_cuda(), + "dpa4c_graph_compress: edge_vec must be a CUDA tensor"); + const c10::cuda::CUDAGuard device_guard(edge_vec.device()); + auto options = edge_vec.options().dtype(torch::kFloat32); + auto descriptor = torch::empty({node_count, widths.output_width}, options); + auto state = torch::empty({node_count, widths.moment_width + 2}, options); + if (node_count == 0) { + return {descriptor, state}; + } + auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous(); + Arguments arguments = + build_arguments(payload, edge_vec_float, channels, widths); + arguments.descriptor = descriptor.data_ptr(); + arguments.state_out = state.data_ptr(); + dispatch(channels, false, arguments, at::cuda::getCurrentCUDAStream()); + return {descriptor, state}; +} + +torch::Tensor dpa4c_graph_compress_backward_impl( + torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor destination_order, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + bool canonical, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor, + bool reuse_state) { + const Payload payload{edge_index, edge_mask, destination_order, + destination_row_ptr, atype, table, + pair_film, pair_mixing, type_embedding, + readout_matrices, coupling_meta, coupling_entry, + coupling_value, output_mean, output_inv_std, + canonical, lmax, table_stride, + table_max, rcut, eps, + degree_floor}; + const long node_count = destination_row_ptr.numel() - 1; + // The destination row pointer defines the node axis; the type table must + // describe exactly that axis, or the two disagree on how many nodes exist. + TORCH_CHECK( + atype.size(0) == destination_row_ptr.numel() - 1, + "dpa4c_graph_compress_backward: atype and destination_row_ptr describe " + "different node counts"); + const int channels = static_cast(type_embedding.size(1)); + const Dimensions widths = + profile_dimensions(channels, static_cast(lmax)); + TORCH_CHECK(edge_vec.is_cuda(), + "dpa4c_graph_compress_backward: edge_vec must be a CUDA tensor"); + const c10::cuda::CUDAGuard device_guard(edge_vec.device()); + TORCH_CHECK(state.is_cuda() && state.device() == edge_vec.device() && + state.is_contiguous() && + state.scalar_type() == torch::kFloat32 && + state.sizes() == + torch::IntArrayRef({node_count, widths.moment_width + 2}), + "dpa4c_graph_compress_backward: invalid saved state"); + TORCH_CHECK( + descriptor_gradient.is_cuda() && + descriptor_gradient.device() == edge_vec.device() && + descriptor_gradient.numel() == node_count * widths.output_width, + "dpa4c_graph_compress_backward: invalid descriptor gradient"); + if (node_count == 0) { + return torch::zeros_like(edge_vec); + } + auto descriptor_gradient_float = + descriptor_gradient.to(torch::kFloat32).contiguous(); + auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous(); + auto edge_gradient = torch::empty_like(edge_vec_float); + // The moment cotangent has exactly the layout of the saved state, so an + // inference caller that no longer needs the state can reuse its storage. + auto moment_gradient = reuse_state ? state : torch::empty_like(state); + Arguments arguments = + build_arguments(payload, edge_vec_float, channels, widths); + arguments.descriptor_gradient = descriptor_gradient_float.data_ptr(); + arguments.state = state.data_ptr(); + arguments.moment_gradient = moment_gradient.data_ptr(); + arguments.edge_gradient = edge_gradient.data_ptr(); + dispatch(channels, true, arguments, at::cuda::getCurrentCUDAStream()); + return edge_gradient.to(edge_vec.scalar_type()); +} + +torch::Tensor dpa4c_graph_compress_backward(torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor destination_order, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + bool canonical, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor) { + return dpa4c_graph_compress_backward_impl( + descriptor_gradient, state, edge_vec, edge_index, edge_mask, + destination_order, destination_row_ptr, atype, table, pair_film, + pair_mixing, type_embedding, readout_matrices, coupling_meta, + coupling_entry, coupling_value, output_mean, output_inv_std, canonical, + lmax, table_stride, table_max, rcut, eps, degree_floor, false); +} + +std::tuple dpa4c_canonical_compress( + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor) { + TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), + "dpa4c_canonical_compress: source and edge_vec must share the " + "edge axis"); + auto edge_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); + auto destination_order = torch::empty({0}, source.options()); + return dpa4c_graph_compress(edge_vec, source, edge_mask, destination_order, + destination_row_ptr, atype, table, pair_film, + pair_mixing, type_embedding, readout_matrices, + coupling_meta, coupling_entry, coupling_value, + output_mean, output_inv_std, true, lmax, + table_stride, table_max, rcut, eps, degree_floor); +} + +torch::Tensor dpa4c_canonical_compress_backward_common( + torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor, + bool reuse_state) { + TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), + "dpa4c_canonical_compress_backward: source and edge_vec must " + "share the edge axis"); + auto edge_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); + auto destination_order = torch::empty({0}, source.options()); + return dpa4c_graph_compress_backward_impl( + descriptor_gradient, state, edge_vec, source, edge_mask, + destination_order, destination_row_ptr, atype, table, pair_film, + pair_mixing, type_embedding, readout_matrices, coupling_meta, + coupling_entry, coupling_value, output_mean, output_inv_std, true, lmax, + table_stride, table_max, rcut, eps, degree_floor, reuse_state); +} + +torch::Tensor dpa4c_canonical_compress_backward( + torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor) { + return dpa4c_canonical_compress_backward_common( + descriptor_gradient, state, edge_vec, source, destination_row_ptr, atype, + table, pair_film, pair_mixing, type_embedding, readout_matrices, + coupling_meta, coupling_entry, coupling_value, output_mean, + output_inv_std, lmax, table_stride, table_max, rcut, eps, degree_floor, + false); +} + +torch::Tensor dpa4c_canonical_compress_backward_inplace( + torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor) { + return dpa4c_canonical_compress_backward_common( + descriptor_gradient, state, edge_vec, source, destination_row_ptr, atype, + table, pair_film, pair_mixing, type_embedding, readout_matrices, + coupling_meta, coupling_entry, coupling_value, output_mean, + output_inv_std, lmax, table_stride, table_max, rcut, eps, degree_floor, + true); +} + +// Energy and edge cotangent of one compressed inference step, evaluated over +// runs of consecutive destination nodes. +// +// Destination-sorted CSR gives a run a contiguous span of the edge axis and +// assigns every edge to exactly one destination, so the runs partition the +// work rather than splitting any reduction. Folding the descriptor, the +// fitting and the descriptor backward into one operator lets a run retire its +// descriptor, its cotangent and its moment state before the next run starts, +// which leaves only the graph and the edge cotangent at system scale. Nothing +// is recomputed: the fitting seed is the ownership mask, known up front. +// +// The loop lives here rather than in Python because its trip count follows a +// dynamic node count, which export cannot trace. +std::tuple +dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor, + std::vector ws, + std::vector bs, + std::vector resnets, + torch::Tensor w_head, + torch::Tensor b_head, + torch::Tensor bias_atom_e, + int64_t act, + torch::Tensor seed, + int64_t tile) { + TORCH_CHECK(edge_vec.is_cuda(), + "dpa4c_canonical_compress_energy_gradient: edge_vec must be a " + "CUDA tensor"); + TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), + "dpa4c_canonical_compress_energy_gradient: source and edge_vec " + "must share the edge axis"); + const c10::cuda::CUDAGuard device_guard(edge_vec.device()); + const long node_count = atype.size(0); + TORCH_CHECK(destination_row_ptr.numel() == node_count + 1, + "dpa4c_canonical_compress_energy_gradient: atype and " + "destination_row_ptr describe different node counts"); + const int channels = static_cast(type_embedding.size(1)); + const Dimensions widths = + profile_dimensions(channels, static_cast(lmax)); + auto f32 = edge_vec.options().dtype(torch::kFloat32); + auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous(); + auto energy = + torch::empty({node_count, 1}, edge_vec.options().dtype(torch::kFloat64)); + // With no destination rows every edge slot is padding, and the loop that + // would clear it never runs. + auto edge_gradient = node_count == 0 ? torch::zeros_like(edge_vec_float) + : torch::empty_like(edge_vec_float); + if (node_count == 0) { + return {energy, edge_gradient.to(edge_vec.scalar_type())}; + } + auto seed_c = seed.contiguous(); + TORCH_CHECK( + seed_c.numel() == node_count && seed_c.scalar_type() == torch::kFloat64, + "dpa4c_canonical_compress_energy_gradient: seed must be fp64 " + "with one entry per node"); + + const FittingLayerPlan plan = fitting_layer_plan(ws); + TORCH_CHECK(plan.n_layer > 0 && ws[0].size(0) == widths.output_width, + "dpa4c_canonical_compress_energy_gradient: first fitting weight " + "does not match the descriptor width"); + const long run = tile > 0 + ? std::max(1, std::min(tile, node_count)) + : node_count; + const int slots = plan.n_layer > 1 ? 2 : 1; + auto descriptor = torch::empty({run, widths.output_width}, f32); + auto state = torch::empty({run, widths.moment_width + 2}, f32); + auto saved = torch::empty({run * plan.saved_width()}, f32); + auto scratch = torch::empty({slots, run, plan.width_max}, f32); + float* slot[2] = { + scratch[0].data_ptr(), + slots > 1 ? scratch[1].data_ptr() : scratch[0].data_ptr()}; + + auto empty_index = torch::empty({0}, source.options()); + auto empty_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); + auto stream = at::cuda::getCurrentCUDAStream(); + for (long begin = 0; begin < node_count; begin += run) { + const long count = std::min(run, node_count - begin); + const Payload payload{ + source, + empty_mask, + empty_index, + destination_row_ptr.slice(0, begin, begin + count + 1), + atype, + table, + pair_film, + pair_mixing, + type_embedding, + readout_matrices, + coupling_meta, + coupling_entry, + coupling_value, + output_mean, + output_inv_std, + true, + lmax, + table_stride, + table_max, + rcut, + eps, + degree_floor, + begin}; + Arguments arguments = + build_arguments(payload, edge_vec_float, channels, widths); + arguments.descriptor = descriptor.data_ptr(); + arguments.state_out = state.data_ptr(); + dispatch(channels, false, arguments, stream); + + fitting_forward_range(stream, plan, descriptor.data_ptr(), + widths.output_width, atype.data_ptr() + begin, + ws, bs, resnets, w_head, b_head, bias_atom_e, act, + count, saved.data_ptr(), slot, + energy.data_ptr() + begin); + // The cotangent replaces the descriptor, which the run no longer needs. + fitting_backward_range( + stream, plan, seed_c.data_ptr() + begin, + saved.data_ptr(), ws, bs, resnets, w_head, act, count, slot[0], + plan.n_layer > 1 ? slot[1] : nullptr, descriptor.data_ptr()); + + arguments.descriptor_gradient = descriptor.data_ptr(); + arguments.state = state.data_ptr(); + arguments.moment_gradient = state.data_ptr(); + arguments.edge_gradient = edge_gradient.data_ptr(); + // Only the final run reaches the reserved edge slots; its row pointer ends + // at the last physical edge, which is exactly where the padding begins. + arguments.clear_padding = begin + count == node_count; + dispatch(channels, true, arguments, stream); + } + return {energy, edge_gradient.to(edge_vec.scalar_type())}; +} + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.def( + "dpa4c_graph_compress(Tensor edge_vec, Tensor edge_index, " + "Tensor edge_mask, Tensor destination_order, " + "Tensor destination_row_ptr, Tensor atype, Tensor table, " + "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " + "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " + "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "bool canonical, int lmax, float table_stride, float table_max, " + "float rcut, float eps, float degree_floor) " + "-> (Tensor descriptor, Tensor state)"); + library.impl("dpa4c_graph_compress", torch::kCUDA, &dpa4c_graph_compress); + library.def( + "dpa4c_graph_compress_backward(Tensor descriptor_gradient, " + "Tensor state, Tensor edge_vec, Tensor edge_index, Tensor edge_mask, " + "Tensor destination_order, Tensor destination_row_ptr, Tensor atype, " + "Tensor table, Tensor pair_film, Tensor pair_mixing, " + "Tensor type_embedding, Tensor readout_matrices, Tensor coupling_meta, " + "Tensor coupling_entry, Tensor coupling_value, Tensor output_mean, " + "Tensor output_inv_std, bool canonical, int lmax, float table_stride, " + "float table_max, float rcut, float eps, float degree_floor) -> Tensor"); + library.impl("dpa4c_graph_compress_backward", torch::kCUDA, + &dpa4c_graph_compress_backward); + library.def( + "dpa4c_canonical_compress(Tensor edge_vec, Tensor source, " + "Tensor destination_row_ptr, Tensor atype, Tensor table, " + "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " + "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " + "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "int lmax, float table_stride, float table_max, float rcut, float eps, " + "float degree_floor) -> (Tensor descriptor, Tensor state)"); + library.impl("dpa4c_canonical_compress", torch::kCUDA, + &dpa4c_canonical_compress); + library.def( + "dpa4c_canonical_compress_backward(Tensor descriptor_gradient, " + "Tensor state, Tensor edge_vec, Tensor source, " + "Tensor destination_row_ptr, Tensor atype, Tensor table, " + "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " + "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " + "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "int lmax, float table_stride, float table_max, float rcut, float eps, " + "float degree_floor) -> Tensor"); + library.impl("dpa4c_canonical_compress_backward", torch::kCUDA, + &dpa4c_canonical_compress_backward); + library.def( + "dpa4c_canonical_compress_backward_inplace(" + "Tensor descriptor_gradient, Tensor(a!) state, Tensor edge_vec, " + "Tensor source, Tensor destination_row_ptr, Tensor atype, Tensor table, " + "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " + "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " + "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "int lmax, float table_stride, float table_max, float rcut, float eps, " + "float degree_floor) -> Tensor"); + library.impl("dpa4c_canonical_compress_backward_inplace", torch::kCUDA, + &dpa4c_canonical_compress_backward_inplace); + library.def( + "dpa4c_canonical_compress_energy_gradient(Tensor edge_vec, " + "Tensor source, Tensor destination_row_ptr, Tensor atype, Tensor table, " + "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " + "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " + "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "int lmax, float table_stride, float table_max, float rcut, float eps, " + "float degree_floor, Tensor[] ws, Tensor[] bs, int[] resnets, " + "Tensor w_head, Tensor b_head, Tensor bias_atom_e, int act, " + "Tensor seed, int tile) -> (Tensor energy, Tensor edge_gradient)"); + library.impl("dpa4c_canonical_compress_energy_gradient", torch::kCUDA, + &dpa4c_canonical_compress_energy_gradient); +} diff --git a/source/op/pt/dpa4c_graph_compress.cuh b/source/op/pt/dpa4c_graph_compress.cuh new file mode 100644 index 0000000000..7daa085a27 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress.cuh @@ -0,0 +1,635 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Device templates for the compressed degree-wise DPA4C descriptor. +// +// One warp owns one destination node. The forward scan evaluates the +// tabulated radial branch and its shared mode profiles, applies the ordered +// PairFiLM amplitude, and reduces both envelope masses together with every +// degree-wise moment in a single pass over the destination CSR row. The +// analytical backward runs one node-readout VJP followed by one edge +// recomputation scan that reproduces the amplitude instead of storing it. +// +// Degrees one and two carry the wide channel blocks and are contracted in +// closed form. Degrees three and four carry a single channel each, so their +// couplings run through a compact sparse Cartesian Gaunt table supplied as a +// compression artifact. + +#pragma once + +#include + +#include + +#include "dpa4c_graph_compress_launch.h" + +namespace deepmd_dpa4c { + +constexpr unsigned kWarpMask = 0xffffffffu; +constexpr int kMaxRadialModes = 8; +constexpr float kSqrtTwo = 1.4142135623730950488f; +constexpr float kSqrtThree = 1.7320508075688772935f; +constexpr float kInvSqrtTwo = 0.7071067811865475244f; +constexpr float kInvSqrtFive = 0.44721359549995793928f; +constexpr float kInvSqrtSix = 0.4082482904638630164f; + +// Unit-Frobenius normalization of the Cartesian 222 Gaunt tensor. The triple +// product needs only one operator ordering: the three factors are symmetric, +// so `tr(ABC) = tr((ABC)^T) = tr(CBA) = tr(ACB)` and the two orderings of the +// contraction are identical. +constexpr float kBis222Scale = -0.58554004376911988f; + +// === Warp primitives === + +__device__ __forceinline__ float warp_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(kWarpMask, value, offset); + } + return __shfl_sync(kWarpMask, value, 0); +} + +template +__device__ __forceinline__ unsigned subwarp_mask(int leader) { + if constexpr (Width == kWarpSize) { + return kWarpMask; + } else { + return ((1u << Width) - 1u) << leader; + } +} + +template +__device__ __forceinline__ float reduce_channel_groups(float value) { + if constexpr (Width < kWarpSize) { +#pragma unroll + for (int offset = Width; offset < kWarpSize; offset <<= 1) { + value += __shfl_xor_sync(kWarpMask, value, offset); + } + } + return value; +} + +template +__device__ __forceinline__ float subwarp_sum(float value, unsigned mask) { +#pragma unroll + for (int offset = Width / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(mask, value, offset, Width); + } + return value; +} + +template +__device__ __forceinline__ long edge_at_position( + long position, const index_t* destination_order) { + if constexpr (Canonical) { + return position; + } + return static_cast(destination_order[position]); +} + +// === Radial table === + +struct TableLocation { + int index; + float coordinate; + bool clamped; +}; + +template +__device__ __forceinline__ TableLocation +locate_table(float radius, float stride, float table_max, int interval_count) { + const float coordinate = + Canonical ? radius : fminf(fmaxf(radius, 0.0f), table_max); + int index = static_cast(__fdividef(coordinate, stride)); + index = min(index, interval_count - 1); + return {index, coordinate - static_cast(index) * stride, + !Canonical && radius >= table_max}; +} + +/// Address of the two coefficient blocks of one channel. +/// +/// The interval row splits into a quartet block and a pair block, so the six +/// spline coefficients arrive in one 128-bit and one 64-bit load rather than +/// three 64-bit loads, at identical traffic. Both bases are resolved once per +/// edge, which turns every channel offset into a compile-time immediate. +struct TableRow { + const float4* quartet; + const float2* pair; +}; + +__device__ __forceinline__ TableRow table_row(const float* table, + const TableLocation& location, + int width) { + const float* row = table + static_cast(location.index) * width * 6; + return {reinterpret_cast(row), + reinterpret_cast(row + 4 * width)}; +} + +__device__ __forceinline__ float evaluate_table(const TableRow& row, + int channel, + float x) { + const float4 low = __ldg(row.quartet + channel); + const float2 high = __ldg(row.pair + channel); + return low.x + + (low.y + (low.z + (low.w + (high.x + high.y * x) * x) * x) * x) * x; +} + +/// Evaluate the spline and its derivative as two independent Horner chains. +/// +/// A simultaneous sweep that carries the derivative alongside the value issues +/// fewer instructions but serializes them into one chain of ten dependent +/// products; the two shorter chains below overlap and measure faster at the +/// widest channel profile. +__device__ __forceinline__ float2 evaluate_table_with_derivative( + const TableRow& row, int channel, float x, bool clamped) { + const float4 low = __ldg(row.quartet + channel); + const float2 high = __ldg(row.pair + channel); + const float value = + low.x + (low.y + (low.z + (low.w + (high.x + high.y * x) * x) * x) * x) * x; + const float derivative = + low.y + (2.0f * low.z + + (3.0f * low.w + (4.0f * high.x + 5.0f * high.y * x) * x) * x) * + x; + return make_float2(value, clamped ? 0.0f : derivative); +} + +/// Accumulate the pair-conditioned radial mode residual of one channel. +/// +/// The mode axis is innermost in the ordered mixing cache, so a supported rank +/// of two, four, or eight is covered by at most one 128-bit and one 64-bit +/// global load, and the shared profile in one matching vector read. Splitting +/// the shared read into pairs or single elements lowers register pressure but +/// measures slower at the widest profiles, where the shared-memory +/// instruction count dominates. +__device__ __forceinline__ void accumulate_modes(const float* mixing, + const float* profile, + int rank, + float& value) { + int mode = 0; + for (; mode + 4 <= rank; mode += 4) { + const float4 weight = __ldg(reinterpret_cast(mixing + mode)); + const float4 shape = *reinterpret_cast(profile + mode); + value = fmaf(weight.x, shape.x, value); + value = fmaf(weight.y, shape.y, value); + value = fmaf(weight.z, shape.z, value); + value = fmaf(weight.w, shape.w, value); + } + if (mode < rank) { + const float2 weight = __ldg(reinterpret_cast(mixing + mode)); + const float2 shape = *reinterpret_cast(profile + mode); + value = fmaf(weight.x, shape.x, value); + value = fmaf(weight.y, shape.y, value); + } +} + +/// Accumulate the mode residual and its distance derivative together. +__device__ __forceinline__ void accumulate_modes_with_derivative( + const float* mixing, + const float* profile, + const float* slope, + int rank, + float& value, + float& derivative) { + int mode = 0; + for (; mode + 4 <= rank; mode += 4) { + const float4 weight = __ldg(reinterpret_cast(mixing + mode)); + const float4 shape = *reinterpret_cast(profile + mode); + const float4 rate = *reinterpret_cast(slope + mode); + value = fmaf(weight.x, shape.x, value); + value = fmaf(weight.y, shape.y, value); + value = fmaf(weight.z, shape.z, value); + value = fmaf(weight.w, shape.w, value); + derivative = fmaf(weight.x, rate.x, derivative); + derivative = fmaf(weight.y, rate.y, derivative); + derivative = fmaf(weight.z, rate.z, derivative); + derivative = fmaf(weight.w, rate.w, derivative); + } + if (mode < rank) { + const float2 weight = __ldg(reinterpret_cast(mixing + mode)); + const float2 shape = *reinterpret_cast(profile + mode); + const float2 rate = *reinterpret_cast(slope + mode); + value = fmaf(weight.x, shape.x, value); + value = fmaf(weight.y, shape.y, value); + derivative = fmaf(weight.x, rate.x, derivative); + derivative = fmaf(weight.y, rate.y, derivative); + } +} + +__device__ __forceinline__ float c3_envelope(float radius, float rcut) { + const float u = fminf(fmaxf(__fdividef(rcut - radius, rcut), 0.0f), 1.0f); + const float x = 1.0f - u; + const float series = + 1.0f + x * (4.0f + x * (10.0f + x * (20.0f + 35.0f * x))); + const float u2 = u * u; + return u2 * u2 * series; +} + +__device__ __forceinline__ float c3_envelope_derivative(float radius, + float rcut) { + if (radius <= 0.0f || radius >= rcut) { + return 0.0f; + } + const float x = __fdividef(radius, rcut); + const float u = 1.0f - x; + const float x2 = x * x; + const float x3 = x2 * x; + const float series = + 1.0f + 4.0f * x + 10.0f * x2 + 20.0f * x3 + 35.0f * x2 * x2; + const float derivative = 4.0f + 20.0f * x + 60.0f * x2 + 140.0f * x3; + const float u2 = u * u; + return __fdividef(-4.0f * u2 * u * series + u2 * u2 * derivative, rcut); +} + +// === Edge geometry and angular basis === + +struct EdgeGeometry { + float ux; + float uy; + float uz; + float radius; + float inverse_radius; + float envelope; + int source_type; +}; + +template +__device__ __forceinline__ EdgeGeometry load_geometry(long edge, + float rcut, + float eps, + const float* edge_vec, + const index_t* edge_index, + const long* atype) { + EdgeGeometry geometry; + const long source = static_cast(edge_index[edge]); + geometry.source_type = static_cast(atype[source]); + const float x = edge_vec[edge * 3 + 0]; + const float y = edge_vec[edge * 3 + 1]; + const float z = edge_vec[edge * 3 + 2]; + const float square = x * x + y * y + z * z + eps * eps; + geometry.inverse_radius = rsqrtf(square); + geometry.radius = square * geometry.inverse_radius; + geometry.ux = x * geometry.inverse_radius; + geometry.uy = y * geometry.inverse_radius; + geometry.uz = z * geometry.inverse_radius; + geometry.envelope = c3_envelope(geometry.radius, rcut); + return geometry; +} + +// The Cartesian harmonics are evaluated on the unit direction, so the squared +// norm that makes each degree traceless is exactly one. Two polynomials that +// agree on the unit sphere differ by a multiple of `|u|^2 - 1`, whose gradient +// at a unit vector is purely radial and is therefore annihilated by the +// tangential projection that closes the coordinate VJP. Substituting the +// constant is thus exact on the unit sphere for both the value and the +// projected gradient. The regularized direction departs from unit norm by a +// relative `eps^2 / rho^2`, far below single precision at any physical +// separation. + +// Degrees zero through two, which every profile evaluates. +__device__ __forceinline__ void fill_angular_basis(const EdgeGeometry& geometry, + float (&basis)[9]) { + basis[0] = 1.0f; + basis[1] = geometry.ux; + basis[2] = geometry.uy; + basis[3] = geometry.uz; + basis[4] = kSqrtThree * geometry.ux * geometry.uy; + basis[5] = kSqrtThree * geometry.uy * geometry.uz; + basis[6] = 0.5f * (3.0f * geometry.uz * geometry.uz - 1.0f); + basis[7] = kSqrtThree * geometry.ux * geometry.uz; + basis[8] = 0.5f * kSqrtThree * + (geometry.ux * geometry.ux - geometry.uy * geometry.uy); +} + +// Real Cartesian harmonics of degrees three and four, addressed by the flat +// index `m` for degree three and `7 + m` for degree four. These degrees carry +// one channel, so their components are distributed across the lanes of an edge +// group. The caller therefore iterates the component at compile time and +// selects with a lane predicate, which folds the selector below into the one +// case that lane needs. +__device__ __forceinline__ float high_basis_value(const EdgeGeometry& geometry, + int index) { + const float x = geometry.ux; + const float y = geometry.uy; + const float z = geometry.uz; + const float z2 = z * z; + const float difference = x * x - y * y; + switch (index) { + case 0: + return 0.79056941504209483f * y * (3.0f * x * x - y * y); + case 1: + return 3.87298334620741689f * x * y * z; + case 2: + return 0.61237243569579452f * y * (5.0f * z2 - 1.0f); + case 3: + return 0.5f * z * (5.0f * z2 - 3.0f); + case 4: + return 0.61237243569579452f * x * (5.0f * z2 - 1.0f); + case 5: + return 1.93649167310370844f * z * difference; + case 6: + return 0.79056941504209483f * x * (x * x - 3.0f * y * y); + case 7: + return 2.95803989154980802f * x * y * difference; + case 8: + return 2.09165006633518887f * y * z * (3.0f * x * x - y * y); + case 9: + return 1.11803398874989485f * x * y * (7.0f * z2 - 1.0f); + case 10: + return 0.79056941504209483f * y * z * (7.0f * z2 - 3.0f); + case 11: + return 0.125f * (35.0f * z2 * z2 - 30.0f * z2 + 3.0f); + case 12: + return 0.79056941504209483f * x * z * (7.0f * z2 - 3.0f); + case 13: + return 0.55901699437494742f * difference * (7.0f * z2 - 1.0f); + case 14: + return 2.09165006633518887f * x * z * (x * x - 3.0f * y * y); + default: + return 0.73950997288745200f * + (x * x * x * x - 6.0f * x * x * y * y + y * y * y * y); + } +} + +// Accumulate `weight * grad B_index` into the Cartesian derivative triple. +__device__ __forceinline__ void high_basis_gradient( + const EdgeGeometry& geometry, int index, float weight, float (&du)[3]) { + const float x = geometry.ux; + const float y = geometry.uy; + const float z = geometry.uz; + const float z2 = z * z; + const float difference = x * x - y * y; + float dx = 0.0f; + float dy = 0.0f; + float dz = 0.0f; + switch (index) { + case 0: { + constexpr float k = 0.79056941504209483f; + dx = 6.0f * k * x * y; + dy = 3.0f * k * difference; + break; + } + case 1: { + constexpr float k = 3.87298334620741689f; + dx = k * y * z; + dy = k * x * z; + dz = k * x * y; + break; + } + case 2: { + constexpr float k = 0.61237243569579452f; + dy = k * (5.0f * z2 - 1.0f); + dz = 10.0f * k * y * z; + break; + } + case 3: { + dz = 0.5f * (15.0f * z2 - 3.0f); + break; + } + case 4: { + constexpr float k = 0.61237243569579452f; + dx = k * (5.0f * z2 - 1.0f); + dz = 10.0f * k * x * z; + break; + } + case 5: { + constexpr float k = 1.93649167310370844f; + dx = 2.0f * k * x * z; + dy = -2.0f * k * y * z; + dz = k * difference; + break; + } + case 6: { + constexpr float k = 0.79056941504209483f; + dx = 3.0f * k * difference; + dy = -6.0f * k * x * y; + break; + } + case 7: { + constexpr float k = 2.95803989154980802f; + dx = k * y * (3.0f * x * x - y * y); + dy = k * x * (x * x - 3.0f * y * y); + break; + } + case 8: { + constexpr float k = 2.09165006633518887f; + dx = 6.0f * k * x * y * z; + dy = 3.0f * k * z * difference; + dz = k * y * (3.0f * x * x - y * y); + break; + } + case 9: { + constexpr float k = 1.11803398874989485f; + dx = k * y * (7.0f * z2 - 1.0f); + dy = k * x * (7.0f * z2 - 1.0f); + dz = 14.0f * k * x * y * z; + break; + } + case 10: { + constexpr float k = 0.79056941504209483f; + dy = k * z * (7.0f * z2 - 3.0f); + dz = k * y * (21.0f * z2 - 3.0f); + break; + } + case 11: { + dz = 17.5f * z2 * z - 7.5f * z; + break; + } + case 12: { + constexpr float k = 0.79056941504209483f; + dx = k * z * (7.0f * z2 - 3.0f); + dz = k * x * (21.0f * z2 - 3.0f); + break; + } + case 13: { + constexpr float k = 0.55901699437494742f; + dx = 2.0f * k * x * (7.0f * z2 - 1.0f); + dy = -2.0f * k * y * (7.0f * z2 - 1.0f); + dz = 14.0f * k * difference * z; + break; + } + case 14: { + constexpr float k = 2.09165006633518887f; + dx = 3.0f * k * z * difference; + dy = -6.0f * k * x * y * z; + dz = k * x * (x * x - 3.0f * y * y); + break; + } + default: { + constexpr float k = 0.73950997288745200f; + dx = 4.0f * k * x * (x * x - 3.0f * y * y); + dy = -4.0f * k * y * (3.0f * x * x - y * y); + break; + } + } + du[0] = fmaf(weight, dx, du[0]); + du[1] = fmaf(weight, dy, du[1]); + du[2] = fmaf(weight, dz, du[2]); +} + +// === Symmetric traceless degree-two algebra === + +struct Matrix3 { + float value[3][3]; +}; + +__device__ __forceinline__ Matrix3 packed_to_stf(const float (&packed)[5]) { + Matrix3 matrix; + matrix.value[0][0] = -packed[2] * kInvSqrtSix + packed[4] * kInvSqrtTwo; + matrix.value[1][1] = -packed[2] * kInvSqrtSix - packed[4] * kInvSqrtTwo; + matrix.value[2][2] = 2.0f * packed[2] * kInvSqrtSix; + matrix.value[0][1] = matrix.value[1][0] = packed[0] * kInvSqrtTwo; + matrix.value[1][2] = matrix.value[2][1] = packed[1] * kInvSqrtTwo; + matrix.value[0][2] = matrix.value[2][0] = packed[3] * kInvSqrtTwo; + return matrix; +} + +__device__ __forceinline__ void matrix_vector(const Matrix3& matrix, + const float (&vector)[3], + float (&output)[3]) { +#pragma unroll + for (int row = 0; row < 3; ++row) { + output[row] = 0.0f; +#pragma unroll + for (int column = 0; column < 3; ++column) { + output[row] = + fmaf(matrix.value[row][column], vector[column], output[row]); + } + } +} + +__device__ __forceinline__ Matrix3 matrix_product(const Matrix3& left, + const Matrix3& right) { + Matrix3 output{}; +#pragma unroll + for (int row = 0; row < 3; ++row) { +#pragma unroll + for (int column = 0; column < 3; ++column) { +#pragma unroll + for (int inner = 0; inner < 3; ++inner) { + output.value[row][column] = + fmaf(left.value[row][inner], right.value[inner][column], + output.value[row][column]); + } + } + } + return output; +} + +__device__ __forceinline__ float matrix_trace(const Matrix3& matrix) { + return matrix.value[0][0] + matrix.value[1][1] + matrix.value[2][2]; +} + +__device__ __forceinline__ void matrix_gradient_to_packed(const Matrix3& matrix, + float (&packed)[5]) { + packed[0] = (matrix.value[0][1] + matrix.value[1][0]) * kInvSqrtTwo; + packed[1] = (matrix.value[1][2] + matrix.value[2][1]) * kInvSqrtTwo; + packed[2] = + (-matrix.value[0][0] - matrix.value[1][1] + 2.0f * matrix.value[2][2]) * + kInvSqrtSix; + packed[3] = (matrix.value[0][2] + matrix.value[2][0]) * kInvSqrtTwo; + packed[4] = (matrix.value[0][0] - matrix.value[1][1]) * kInvSqrtTwo; +} + +// === Descriptor and readout helpers === + +__device__ __forceinline__ long gram_pair_position(int first, + int second, + int width) { + const int row = min(first, second); + const int column = max(first, second); + return static_cast(row) * width - + static_cast(row) * (row - 1) / 2 + column - row; +} + +__device__ __forceinline__ void decode_upper_pair(int pair, + int width, + int& row, + int& column) { + row = 0; + while (pair >= width - row) { + pair -= width - row; + ++row; + } + column = row + pair; +} + +__device__ __forceinline__ void store_descriptor(float* descriptor, + const float* mean, + const float* inverse_stddev, + long node, + int output_width, + int coordinate, + float value) { + const long index = node * output_width + coordinate; + descriptor[index] = + (value - __ldg(mean + coordinate)) * __ldg(inverse_stddev + coordinate); +} + +__device__ __forceinline__ float load_output_gradient( + const float* gradient, + const float* inverse_stddev, + long node, + int output_width, + int coordinate) { + return __ldg(gradient + node * output_width + coordinate) * + __ldg(inverse_stddev + coordinate); +} + +template +__device__ __forceinline__ float readout_weight(const float* matrices, + int matrix, + int row, + int column) { + using P = Profile; + return __ldg(matrices + (static_cast(matrix) * P::C1 + row) * P::C1 + + column); +} + +// Probe coordinate of one degree. Degrees three and above carry a single +// channel whose alignment and probe projections are both the identity, so +// their probes are the stored moments. +template +__device__ __forceinline__ float probe_value(const float* probes, + const float* moments, + int degree, + int component, + int rank_index) { + using P = Profile; + if (degree == 1) { + return probes[component * P::K1 + rank_index]; + } + if (degree == 2) { + return probes[3 * P::K1 + component * P::K2 + rank_index]; + } + return moments[P::HighOffset + (degree == 3 ? 0 : P::High3) + component]; +} + +// The Cartesian basis VJP maps angular cotangents to a coordinate gradient. +// Applying it per lane reduces three Cartesian components instead of the full +// set of angular components across the edge group. +__device__ __forceinline__ void basis_vjp(const EdgeGeometry& geometry, + const float (&d_basis)[9], + const float (&high_du)[3], + float radial_gradient, + float (&output)[3]) { + const float dux = + high_du[0] + d_basis[1] + + kSqrtThree * (d_basis[4] * geometry.uy + d_basis[7] * geometry.uz + + d_basis[8] * geometry.ux); + const float duy = + high_du[1] + d_basis[2] + + kSqrtThree * (d_basis[4] * geometry.ux + d_basis[5] * geometry.uz - + d_basis[8] * geometry.uy); + const float duz = + high_du[2] + d_basis[3] + + kSqrtThree * (d_basis[5] * geometry.uy + d_basis[7] * geometry.ux) + + 3.0f * d_basis[6] * geometry.uz; + const float dot = geometry.ux * dux + geometry.uy * duy + geometry.uz * duz; + output[0] = (dux - geometry.ux * dot) * geometry.inverse_radius + + radial_gradient * geometry.ux; + output[1] = (duy - geometry.uy * dot) * geometry.inverse_radius + + radial_gradient * geometry.uy; + output[2] = (duz - geometry.uz * dot) * geometry.inverse_radius + + radial_gradient * geometry.uz; +} + +} // namespace deepmd_dpa4c diff --git a/source/op/pt/dpa4c_graph_compress_c128.cu b/source/op/pt/dpa4c_graph_compress_c128.cu new file mode 100644 index 0000000000..fe39aa3bed --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress_c128.cu @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compiled specializations of the compressed DPA4C descriptor for a scalar +// width of 128 channels. + +#include "dpa4c_graph_compress_kernel.cuh" + +namespace deepmd_dpa4c { + +DPA4C_DEFINE_CHANNEL(128) + +} // namespace deepmd_dpa4c diff --git a/source/op/pt/dpa4c_graph_compress_c16.cu b/source/op/pt/dpa4c_graph_compress_c16.cu new file mode 100644 index 0000000000..50a7a5a88d --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress_c16.cu @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compiled specializations of the compressed DPA4C descriptor for a scalar +// width of 16 channels. + +#include "dpa4c_graph_compress_kernel.cuh" + +namespace deepmd_dpa4c { + +DPA4C_DEFINE_CHANNEL(16) + +} // namespace deepmd_dpa4c diff --git a/source/op/pt/dpa4c_graph_compress_c32.cu b/source/op/pt/dpa4c_graph_compress_c32.cu new file mode 100644 index 0000000000..74fd14c983 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress_c32.cu @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compiled specializations of the compressed DPA4C descriptor for a scalar +// width of 32 channels. + +#include "dpa4c_graph_compress_kernel.cuh" + +namespace deepmd_dpa4c { + +DPA4C_DEFINE_CHANNEL(32) + +} // namespace deepmd_dpa4c diff --git a/source/op/pt/dpa4c_graph_compress_c64.cu b/source/op/pt/dpa4c_graph_compress_c64.cu new file mode 100644 index 0000000000..46d857cc57 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress_c64.cu @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compiled specializations of the compressed DPA4C descriptor for a scalar +// width of 64 channels. + +#include "dpa4c_graph_compress_kernel.cuh" + +namespace deepmd_dpa4c { + +DPA4C_DEFINE_CHANNEL(64) + +} // namespace deepmd_dpa4c diff --git a/source/op/pt/dpa4c_graph_compress_c8.cu b/source/op/pt/dpa4c_graph_compress_c8.cu new file mode 100644 index 0000000000..38428827d1 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress_c8.cu @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Compiled specializations of the compressed DPA4C descriptor for a scalar +// width of 8 channels. + +#include "dpa4c_graph_compress_kernel.cuh" + +namespace deepmd_dpa4c { + +DPA4C_DEFINE_CHANNEL(8) + +} // namespace deepmd_dpa4c diff --git a/source/op/pt/dpa4c_graph_compress_kernel.cuh b/source/op/pt/dpa4c_graph_compress_kernel.cuh new file mode 100644 index 0000000000..c315be1af5 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress_kernel.cuh @@ -0,0 +1,1463 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Kernels of the compressed degree-wise DPA4C descriptor. +// +// This header is included by exactly one translation unit per scalar width, +// which instantiates the angular-degree and topology specializations of that +// width. Keeping the instantiations separated bounds the compile time of any +// single translation unit and lets them proceed in parallel. + +#pragma once + +#include + +#include "dpa4c_graph_compress.cuh" + +namespace deepmd_dpa4c { + +// Shared mode profiles are cached per edge group. The stride is a multiple of +// four so that each row is aligned for the vector reads of the mode residual, +// and it keeps the concurrent groups of one warp on distinct banks. +constexpr int kModeStride = kMaxRadialModes + 4; + +// === Forward === + +template +__global__ __launch_bounds__(Profile::Threads, + 32) void forward_kernel(Arguments args) { + using P = Profile; + constexpr int EdgeWidth = P::ForwardEdgeWidth; + constexpr int Groups = kWarpSize / EdgeWidth; + constexpr int ChannelTiles = Channels / EdgeWidth; + constexpr int AngularTiles = (P::C1 + EdgeWidth - 1) / EdgeWidth; + constexpr int TensorTiles = (P::C2 + EdgeWidth - 1) / EdgeWidth; + constexpr int HighTiles = (P::HighCount + EdgeWidth - 1) / EdgeWidth; + + const int thread = threadIdx.x; + const long node = blockIdx.x; + if (node >= args.node_count) { + return; + } + const int base_channel = thread & (EdgeWidth - 1); + const int group = thread / EdgeWidth; + const int leader = group * EdgeWidth; + const unsigned mask = subwarp_mask(leader); + const auto* edge_index = static_cast(args.edge_index); + const auto* destination_order = + static_cast(args.destination_order); + const int center_type = static_cast(args.atype[args.node_begin + node]); + const long begin = args.destination_row_ptr[node]; + const long end = args.destination_row_ptr[node + 1]; + const int radial_modes = HasModes ? args.radial_modes : 0; + + float scalar[ChannelTiles] = {}; + float vector[AngularTiles][3] = {}; + float tensor[TensorTiles][5] = {}; + float high[HighTiles > 0 ? HighTiles : 1] = {}; + float scalar_mass = 0.0f; + float angular_mass = 0.0f; + + __shared__ float mode_cache[HasModes ? Groups * kModeStride : 1]; + float* modes = mode_cache + (HasModes ? group * kModeStride : 0); + + // === Step 1. Reduce the destination row into degree-wise moments === + for (long position = begin + group; position < end; position += Groups) { + const long edge = edge_at_position(position, destination_order); + if (args.edge_mask != nullptr && !args.edge_mask[edge]) { + continue; + } + // Every lane of the group reloads the shared edge state instead of + // broadcasting it from a leader. The addresses are identical inside the + // group, so the memory system serves one transaction either way, whereas + // a leader-only branch pays the same issue slots and adds ten shuffles. + const EdgeGeometry geometry = load_geometry( + edge, args.rcut, args.eps, args.edge_vec, edge_index, args.atype); + const TableLocation location = + locate_table(geometry.radius, args.table_stride, + args.table_max, args.interval_count); + if constexpr (!Canonical) { + if (center_type >= args.type_count - 1 || + geometry.source_type >= args.type_count - 1) { + continue; + } + } + const TableRow row = table_row(args.table, location, args.table_width); + const float coordinate = location.coordinate; + if constexpr (HasModes) { + __syncwarp(mask); + for (int mode = base_channel; mode < radial_modes; mode += EdgeWidth) { + modes[mode] = evaluate_table(row, Channels + mode, coordinate); + } + __syncwarp(mask); + } + + float basis[9]; + fill_angular_basis(geometry, basis); + const float envelope = geometry.envelope; + const long pair = + static_cast(center_type) * args.type_count + geometry.source_type; + const float2* film_row = + reinterpret_cast(args.pair_film + pair * Channels * 2) + + base_channel; + const float* mixing_row = + HasModes + ? args.pair_mixing + (pair * Channels + base_channel) * radial_modes + : nullptr; + float angular_zero = 0.0f; +#pragma unroll + for (int tile = 0; tile < ChannelTiles; ++tile) { + const int channel = base_channel + tile * EdgeWidth; + const float radial = evaluate_table(row, channel, coordinate); + const float2 film = __ldg(film_row + tile * EdgeWidth); + float film_value = fmaf(film.x, radial, film.y); + if constexpr (HasModes) { + accumulate_modes(mixing_row + tile * (EdgeWidth * radial_modes), modes, + radial_modes, film_value); + } + // Degree zero carries one envelope factor; every non-scalar degree + // carries a second one. + const float amplitude = film_value * envelope; + const float angular = amplitude * envelope; + scalar[tile] += amplitude; + if (channel < P::C1) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + vector[tile][component] = + fmaf(angular, basis[1 + component], vector[tile][component]); + } + } + if (channel < P::C2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + tensor[tile][component] = + fmaf(angular, basis[4 + component], tensor[tile][component]); + } + } + if constexpr (Lmax >= 3) { + if (channel == 0) { + angular_zero = angular; + } + } + } + if constexpr (Lmax >= 3) { + // Degrees three and above read only the leading channel, so their + // components are distributed across the lanes of the edge group. The + // component is iterated at compile time and selected with a lane + // predicate, which folds the harmonic selector and keeps the tile index + // constant. + const float amplitude = __shfl_sync(mask, angular_zero, leader); +#pragma unroll + for (int component = 0; component < P::HighCount; ++component) { + if (component % EdgeWidth == base_channel) { + high[component / EdgeWidth] = + fmaf(amplitude, high_basis_value(geometry, component), + high[component / EdgeWidth]); + } + } + } + if (base_channel == 0) { + const float squared = envelope * envelope; + scalar_mass += squared; + angular_mass = fmaf(squared, squared, angular_mass); + } + } + + // === Step 2. Merge the concurrent edge groups and normalize === + if constexpr (Groups > 1) { +#pragma unroll + for (int tile = 0; tile < ChannelTiles; ++tile) { + scalar[tile] = reduce_channel_groups(scalar[tile]); + } +#pragma unroll + for (int tile = 0; tile < AngularTiles; ++tile) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + vector[tile][component] = + reduce_channel_groups(vector[tile][component]); + } + } +#pragma unroll + for (int tile = 0; tile < TensorTiles; ++tile) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + tensor[tile][component] = + reduce_channel_groups(tensor[tile][component]); + } + } + if constexpr (Lmax >= 3) { +#pragma unroll + for (int tile = 0; tile < HighTiles; ++tile) { + high[tile] = reduce_channel_groups(high[tile]); + } + } + } + __shared__ float normalizer_shared[4]; + { + const float total_scalar = warp_sum(scalar_mass); + const float total_angular = warp_sum(angular_mass); + if (thread == 0) { + normalizer_shared[0] = rsqrtf(total_scalar + args.degree_floor); + normalizer_shared[1] = rsqrtf(total_angular + args.degree_floor); + normalizer_shared[2] = sqrtf(total_scalar + args.degree_floor); + normalizer_shared[3] = sqrtf(total_angular + args.degree_floor); + } + } + __syncthreads(); + const float scalar_norm = normalizer_shared[0]; + const float angular_norm = normalizer_shared[1]; + if (thread < 2) { + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputDivisor + thread, + normalizer_shared[2 + thread]); + } + + // === Step 3. Publish the normalized moments and the saved state === + __shared__ float moments[P::MomentWidth]; + if (thread < EdgeWidth) { +#pragma unroll + for (int tile = 0; tile < ChannelTiles; ++tile) { + moments[P::ScalarOffset + base_channel + tile * EdgeWidth] = + scalar[tile] * scalar_norm; + } +#pragma unroll + for (int tile = 0; tile < AngularTiles; ++tile) { + const int channel = base_channel + tile * EdgeWidth; + if (channel < P::C1) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + moments[P::VectorOffset + component * P::C1 + channel] = + vector[tile][component] * angular_norm; + } + } + } +#pragma unroll + for (int tile = 0; tile < TensorTiles; ++tile) { + const int channel = base_channel + tile * EdgeWidth; + if (channel < P::C2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + moments[P::TensorOffset + component * P::C2 + channel] = + tensor[tile][component] * angular_norm; + } + } + } + if constexpr (Lmax >= 3) { +#pragma unroll + for (int component = 0; component < P::HighCount; ++component) { + if (component % EdgeWidth == base_channel) { + moments[P::HighOffset + component] = + high[component / EdgeWidth] * angular_norm; + } + } + } + } + __syncthreads(); + + const long state_offset = node * P::StateWidth; + for (int coordinate = thread; coordinate < P::MomentWidth; + coordinate += P::Threads) { + args.state_out[state_offset + coordinate] = moments[coordinate]; + } + if (thread == 0) { + args.state_out[state_offset + P::MomentWidth] = scalar_norm; + args.state_out[state_offset + P::MomentWidth + 1] = angular_norm; + } + + // === Step 4. Align and project the wide degrees === + __shared__ float aligned[P::AlignedWidth]; + if (thread < P::C1) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + float value = 0.0f; + for (int input = 0; input < P::C1; ++input) { + value = fmaf(moments[P::VectorOffset + component * P::C1 + input], + readout_weight(args.readout_matrices, 0, + input, thread), + value); + } + aligned[component * P::C1 + thread] = value; + } + } + if (thread < P::C2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + float value = 0.0f; + for (int input = 0; input < P::C2; ++input) { + value = fmaf(moments[P::TensorOffset + component * P::C2 + input], + readout_weight(args.readout_matrices, 2, + input, thread), + value); + } + aligned[3 * P::C1 + component * P::C2 + thread] = value; + } + } + __syncthreads(); + + __shared__ float probes[P::ProbeWidth]; + if (thread < P::K1) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + float value = 0.0f; + for (int input = 0; input < P::C1; ++input) { + value = fmaf(aligned[component * P::C1 + input], + readout_weight(args.readout_matrices, 4, + input, thread), + value); + } + probes[component * P::K1 + thread] = value; + } + } + if (thread < P::K2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + float value = 0.0f; + for (int input = 0; input < P::C2; ++input) { + value = fmaf(aligned[3 * P::C1 + component * P::C2 + input], + readout_weight(args.readout_matrices, 6, + input, thread), + value); + } + probes[3 * P::K1 + component * P::K2 + thread] = value; + } + } + __syncthreads(); + + // === Step 5. Emit the invariant blocks === + for (int channel = thread; channel < Channels; channel += P::Threads) { + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputScalar + channel, + moments[P::ScalarOffset + channel]); + store_descriptor( + args.descriptor, args.output_mean, args.output_inv_std, node, + P::OutputWidth, P::OutputType + channel, + __ldg(args.type_embedding + static_cast(center_type) * Channels + + channel)); + } + + for (int pair = thread; pair < P::Gram1; pair += P::Threads) { + int row, column; + decode_upper_pair(pair, P::C1, row, column); + float value = 0.0f; +#pragma unroll + for (int component = 0; component < 3; ++component) { + value = fmaf(aligned[component * P::C1 + row], + aligned[component * P::C1 + column], value); + } + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputGram1 + pair, + (row == column ? 1.0f : kSqrtTwo) * value); + } + for (int pair = thread; pair < P::Gram2; pair += P::Threads) { + int row, column; + decode_upper_pair(pair, P::C2, row, column); + float value = 0.0f; +#pragma unroll + for (int component = 0; component < 5; ++component) { + value = fmaf(aligned[3 * P::C1 + component * P::C2 + row], + aligned[3 * P::C1 + component * P::C2 + column], value); + } + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputGram2 + pair, + (row == column ? 1.0f : kSqrtTwo) * value); + } + if constexpr (Lmax >= 3) { + if (thread < Lmax - 2) { + const int offset = P::HighOffset + (thread == 0 ? 0 : P::High3); + const int count = 2 * (3 + thread) + 1; + float value = 0.0f; + for (int component = 0; component < count; ++component) { + const float moment = moments[offset + component]; + value = fmaf(moment, moment, value); + } + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputGram3 + thread, value); + } + } + + for (int output = thread; output < P::Bis112; output += P::Threads) { + const int tensor_index = output % P::K2; + int first, second; + decode_upper_pair(output / P::K2, P::K1, first, second); + float packed[5]; + float left[3]; + float right[3]; +#pragma unroll + for (int component = 0; component < 5; ++component) { + packed[component] = probes[3 * P::K1 + component * P::K2 + tensor_index]; + } +#pragma unroll + for (int component = 0; component < 3; ++component) { + left[component] = probes[component * P::K1 + first]; + right[component] = probes[component * P::K1 + second]; + } + float product[3]; + matrix_vector(packed_to_stf(packed), right, product); + const float value = + -kInvSqrtFive * + (left[0] * product[0] + left[1] * product[1] + left[2] * product[2]); + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputBis112 + output, + (first == second ? 1.0f : kSqrtTwo) * value); + } + + if (thread < P::Bis222) { + constexpr int entries[4][3] = {{0, 0, 0}, {0, 0, 1}, {0, 1, 1}, {1, 1, 1}}; + constexpr float scales[4] = {1.0f, kSqrtThree, kSqrtThree, 1.0f}; + Matrix3 matrices[3]; +#pragma unroll + for (int axis = 0; axis < 3; ++axis) { + float packed[5]; +#pragma unroll + for (int component = 0; component < 5; ++component) { + packed[component] = + probes[3 * P::K1 + component * P::K2 + entries[thread][axis]]; + } + matrices[axis] = packed_to_stf(packed); + } + const float value = + kBis222Scale * + matrix_trace(matrix_product(matrix_product(matrices[0], matrices[1]), + matrices[2])); + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputBis222 + thread, + scales[thread] * value); + } + + if constexpr (Lmax >= 3) { + for (int record = 0; record < args.coupling_count; ++record) { + const int* meta = args.coupling_meta + record * 8; + const int degree_1 = meta[0]; + const int degree_2 = meta[1]; + const int degree_3 = meta[2]; + const int nonzero_begin = meta[3]; + const int nonzero_count = meta[4]; + const int probe_begin = meta[5]; + const int probe_count = meta[6]; + const int coordinate = meta[7]; + for (int output = thread; output < probe_count; output += P::Threads) { + const int selection = __ldg(args.coupling_entry + probe_begin + output); + const int index_1 = selection & 0xFF; + const int index_2 = (selection >> 8) & 0xFF; + const int index_3 = (selection >> 16) & 0xFF; + float value = 0.0f; + for (int term = 0; term < nonzero_count; ++term) { + const int components = + __ldg(args.coupling_entry + nonzero_begin + term); + const float weight = + __ldg(args.coupling_value + nonzero_begin + term); + const float first = probe_value( + probes, moments, degree_1, components & 0xFF, index_1); + const float second = probe_value( + probes, moments, degree_2, (components >> 8) & 0xFF, index_2); + const float third = probe_value( + probes, moments, degree_3, (components >> 16) & 0xFF, index_3); + value = fmaf(weight * first * second, third, value); + } + store_descriptor( + args.descriptor, args.output_mean, args.output_inv_std, node, + P::OutputWidth, coordinate + output, + value * __ldg(args.coupling_value + probe_begin + output)); + } + } + } + + for (int output = thread; output < P::Quartic; output += P::Threads) { + const int vector_index = output % P::K1; + const int tensor_index = output / P::K1; + float packed[5]; + float value[3]; +#pragma unroll + for (int component = 0; component < 5; ++component) { + packed[component] = probes[3 * P::K1 + component * P::K2 + tensor_index]; + } +#pragma unroll + for (int component = 0; component < 3; ++component) { + value[component] = probes[component * P::K1 + vector_index]; + } + float product[3]; + matrix_vector(packed_to_stf(packed), value, product); + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputQuartic + output, + product[0] * product[0] + product[1] * product[1] + + product[2] * product[2]); + } +} + +// === Node readout backward === + +// The four symmetric 222 outputs are +// k * {tr(Q0^3), sqrt(3) tr(Q0^2 Q1), sqrt(3) tr(Q0 Q1^2), tr(Q1^3)}. +// Their gradients need only Q0^2, Q1^2, and the symmetrized Q0 Q1 product, so +// evaluating this closed form inside the node backward avoids a second probe +// projection and the associated global gradient checkpoint. +template +__device__ __forceinline__ void add_bis222_probe_gradient( + int lane, + long node, + const Arguments& args, + const float* __restrict__ probes, + float (&d_tensor)[5]) { + using P = Profile; + float packed_0[5]; + float packed_1[5]; +#pragma unroll + for (int component = 0; component < 5; ++component) { + packed_0[component] = probes[component * P::K2 + 0]; + packed_1[component] = probes[component * P::K2 + 1]; + } + const Matrix3 matrix_0 = packed_to_stf(packed_0); + const Matrix3 matrix_1 = packed_to_stf(packed_1); + const float gradient_000 = + kBis222Scale * load_output_gradient(args.descriptor_gradient, + args.output_inv_std, node, + P::OutputWidth, P::OutputBis222 + 0); + const float gradient_001 = + kBis222Scale * kSqrtThree * + load_output_gradient(args.descriptor_gradient, args.output_inv_std, node, + P::OutputWidth, P::OutputBis222 + 1); + const float gradient_011 = + kBis222Scale * kSqrtThree * + load_output_gradient(args.descriptor_gradient, args.output_inv_std, node, + P::OutputWidth, P::OutputBis222 + 2); + const float gradient_111 = + kBis222Scale * load_output_gradient(args.descriptor_gradient, + args.output_inv_std, node, + P::OutputWidth, P::OutputBis222 + 3); + + Matrix3 matrix_gradient{}; +#pragma unroll + for (int row = 0; row < 3; ++row) { +#pragma unroll + for (int column = 0; column < 3; ++column) { + float product_00 = 0.0f; + float product_11 = 0.0f; + float product_01 = 0.0f; + float product_10 = 0.0f; +#pragma unroll + for (int inner = 0; inner < 3; ++inner) { + product_00 = fmaf(matrix_0.value[row][inner], + matrix_0.value[inner][column], product_00); + product_11 = fmaf(matrix_1.value[row][inner], + matrix_1.value[inner][column], product_11); + product_01 = fmaf(matrix_0.value[row][inner], + matrix_1.value[inner][column], product_01); + product_10 = fmaf(matrix_1.value[row][inner], + matrix_0.value[inner][column], product_10); + } + matrix_gradient.value[row][column] = + lane == 0 ? 3.0f * gradient_000 * product_00 + + gradient_001 * (product_01 + product_10) + + gradient_011 * product_11 + : gradient_001 * product_00 + + gradient_011 * (product_01 + product_10) + + 3.0f * gradient_111 * product_11; + } + } + float packed_gradient[5]; + matrix_gradient_to_packed(matrix_gradient, packed_gradient); +#pragma unroll + for (int component = 0; component < 5; ++component) { + d_tensor[component] += packed_gradient[component]; + } +} + +// Scatter the sparse-coupling VJP of one probe slot into a shared scratch row. +// The scratch is indexed by harmonic component, which a register array cannot +// address without spilling, and it is private to the lane, so the reduction +// stays deterministic. +template +__device__ __forceinline__ void accumulate_coupling_gradient( + long node, + const Arguments& args, + const float* __restrict__ probes, + const float* __restrict__ moments, + int degree, + int rank_index, + float* __restrict__ scratch) { + using P = Profile; + for (int record = 0; record < args.coupling_count; ++record) { + const int* meta = args.coupling_meta + record * 8; + const int degrees[3] = {meta[0], meta[1], meta[2]}; + if (degrees[0] != degree && degrees[1] != degree && degrees[2] != degree) { + continue; + } + const int nonzero_begin = meta[3]; + const int nonzero_count = meta[4]; + const int probe_begin = meta[5]; + const int probe_count = meta[6]; + const int coordinate = meta[7]; + for (int output = 0; output < probe_count; ++output) { + const int selection = __ldg(args.coupling_entry + probe_begin + output); + const int indices[3] = {selection & 0xFF, (selection >> 8) & 0xFF, + (selection >> 16) & 0xFF}; + const bool active[3] = {degrees[0] == degree && indices[0] == rank_index, + degrees[1] == degree && indices[1] == rank_index, + degrees[2] == degree && indices[2] == rank_index}; + if (!active[0] && !active[1] && !active[2]) { + continue; + } + const float upstream = + __ldg(args.coupling_value + probe_begin + output) * + load_output_gradient(args.descriptor_gradient, args.output_inv_std, + node, P::OutputWidth, coordinate + output); + for (int term = 0; term < nonzero_count; ++term) { + const int components = + __ldg(args.coupling_entry + nonzero_begin + term); + const int component[3] = {components & 0xFF, (components >> 8) & 0xFF, + (components >> 16) & 0xFF}; + const float weight = + upstream * __ldg(args.coupling_value + nonzero_begin + term); + const float first = probe_value( + probes, moments, degrees[0], component[0], indices[0]); + const float second = probe_value( + probes, moments, degrees[1], component[1], indices[1]); + const float third = probe_value( + probes, moments, degrees[2], component[2], indices[2]); + if (active[0]) { + scratch[component[0]] += weight * second * third; + } + if (active[1]) { + scratch[component[1]] += weight * first * third; + } + if (active[2]) { + scratch[component[2]] += weight * first * second; + } + } + } + } +} + +// Four independent lane groups share one warp. An incomplete final block +// aliases inactive groups to the last valid node so every thread reaches each +// block-wide barrier; stores from those groups are suppressed. +template +__global__ __launch_bounds__(Profile::Threads, + 2) void node_backward_kernel(Arguments args) { + using P = Profile; + constexpr int MaxComponents = 9; + const int thread = threadIdx.x; + const int group = thread / P::NodeWidth; + const int lane = thread & (P::NodeWidth - 1); + const long candidate = static_cast(blockIdx.x) * P::NodeGroups + group; + const bool active = candidate < args.node_count; + const long node = active ? candidate : args.node_count - 1; + const long state_offset = node * P::StateWidth; + + __shared__ float moments_storage[P::NodeGroups * P::MomentWidth]; + float* moments = moments_storage + group * P::MomentWidth; + for (int coordinate = lane; coordinate < P::MomentWidth; + coordinate += P::NodeWidth) { + moments[coordinate] = __ldg(args.state + state_offset + coordinate); + } + const float scalar_norm = __ldg(args.state + state_offset + P::MomentWidth); + const float angular_norm = + __ldg(args.state + state_offset + P::MomentWidth + 1); + __syncthreads(); + + __shared__ float aligned_storage[P::NodeGroups * P::AlignedWidth]; + float* aligned = aligned_storage + group * P::AlignedWidth; + for (int channel = lane; channel < P::C1; channel += P::NodeWidth) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + float value = 0.0f; + for (int input = 0; input < P::C1; ++input) { + value = fmaf(moments[P::VectorOffset + component * P::C1 + input], + readout_weight(args.readout_matrices, 0, + input, channel), + value); + } + aligned[component * P::C1 + channel] = value; + } + } + if (lane < P::C2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + float value = 0.0f; + for (int input = 0; input < P::C2; ++input) { + value = fmaf(moments[P::TensorOffset + component * P::C2 + input], + readout_weight(args.readout_matrices, 2, + input, lane), + value); + } + aligned[3 * P::C1 + component * P::C2 + lane] = value; + } + } + __syncthreads(); + + __shared__ float probes_storage[P::NodeGroups * P::ProbeWidth]; + float* probes = probes_storage + group * P::ProbeWidth; + if (lane < P::K1) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + float value = 0.0f; + for (int input = 0; input < P::C1; ++input) { + value = fmaf(aligned[component * P::C1 + input], + readout_weight(args.readout_matrices, 4, + input, lane), + value); + } + probes[component * P::K1 + lane] = value; + } + } + if (lane < P::K2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + float value = 0.0f; + for (int input = 0; input < P::C2; ++input) { + value = fmaf(aligned[3 * P::C1 + component * P::C2 + input], + readout_weight(args.readout_matrices, 6, + input, lane), + value); + } + probes[3 * P::K1 + component * P::K2 + lane] = value; + } + } + __syncthreads(); + + // Sparse-coupling scratch, one private row per lane. + __shared__ float + coupling_scratch[Lmax >= 3 ? P::NodeGroups * P::NodeWidth * MaxComponents + : 1]; + float* scratch = nullptr; + if constexpr (Lmax >= 3) { + scratch = coupling_scratch + (group * P::NodeWidth + lane) * MaxComponents; +#pragma unroll + for (int component = 0; component < MaxComponents; ++component) { + scratch[component] = 0.0f; + } + } + + __shared__ float d_probes_storage[P::NodeGroups * P::ProbeWidth]; + float* d_probes = d_probes_storage + group * P::ProbeWidth; + if (lane < P::K1) { + float d_vector[3] = {}; + for (int output = 0; output < P::Bis112; ++output) { + const int tensor_index = output % P::K2; + int first, second; + decode_upper_pair(output / P::K2, P::K1, first, second); + if (lane != first && lane != second) { + continue; + } + float packed[5]; + float other[3]; +#pragma unroll + for (int component = 0; component < 5; ++component) { + packed[component] = + probes[3 * P::K1 + component * P::K2 + tensor_index]; + } + const int other_index = lane == first ? second : first; +#pragma unroll + for (int component = 0; component < 3; ++component) { + other[component] = probes[component * P::K1 + other_index]; + } + const float gradient = + -kInvSqrtFive * (first == second ? 2.0f : kSqrtTwo) * + load_output_gradient(args.descriptor_gradient, args.output_inv_std, + node, P::OutputWidth, P::OutputBis112 + output); + float product[3]; + matrix_vector(packed_to_stf(packed), other, product); +#pragma unroll + for (int component = 0; component < 3; ++component) { + d_vector[component] = + fmaf(gradient, product[component], d_vector[component]); + } + } + for (int tensor_index = 0; tensor_index < P::K2; ++tensor_index) { + const float gradient = load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, P::OutputWidth, + P::OutputQuartic + tensor_index * P::K1 + lane); + float packed[5]; + float value[3]; +#pragma unroll + for (int component = 0; component < 5; ++component) { + packed[component] = + probes[3 * P::K1 + component * P::K2 + tensor_index]; + } +#pragma unroll + for (int component = 0; component < 3; ++component) { + value[component] = probes[component * P::K1 + lane]; + } + const Matrix3 matrix = packed_to_stf(packed); + float first[3]; + float second[3]; + matrix_vector(matrix, value, first); + matrix_vector(matrix, first, second); +#pragma unroll + for (int component = 0; component < 3; ++component) { + d_vector[component] = + fmaf(2.0f * gradient, second[component], d_vector[component]); + } + } + if constexpr (Lmax >= 3) { + accumulate_coupling_gradient(node, args, probes, moments, + 1, lane, scratch); +#pragma unroll + for (int component = 0; component < 3; ++component) { + d_vector[component] += scratch[component]; + scratch[component] = 0.0f; + } + } +#pragma unroll + for (int component = 0; component < 3; ++component) { + d_probes[component * P::K1 + lane] = d_vector[component]; + } + } + + if (lane < P::K2) { + float d_tensor[5] = {}; + add_bis222_probe_gradient(lane, node, args, + probes + 3 * P::K1, d_tensor); + for (int output = lane; output < P::Bis112; output += P::K2) { + int first, second; + decode_upper_pair(output / P::K2, P::K1, first, second); + float left[3]; + float right[3]; +#pragma unroll + for (int component = 0; component < 3; ++component) { + left[component] = probes[component * P::K1 + first]; + right[component] = probes[component * P::K1 + second]; + } + const float gradient = + -kInvSqrtFive * (first == second ? 1.0f : kSqrtTwo) * + load_output_gradient(args.descriptor_gradient, args.output_inv_std, + node, P::OutputWidth, P::OutputBis112 + output); + Matrix3 matrix{}; +#pragma unroll + for (int row = 0; row < 3; ++row) { +#pragma unroll + for (int column = 0; column < 3; ++column) { + matrix.value[row][column] = gradient * left[row] * right[column]; + } + } + float packed_gradient[5]; + matrix_gradient_to_packed(matrix, packed_gradient); +#pragma unroll + for (int component = 0; component < 5; ++component) { + d_tensor[component] += packed_gradient[component]; + } + } + for (int vector_index = 0; vector_index < P::K1; ++vector_index) { + const float gradient = load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, P::OutputWidth, + P::OutputQuartic + lane * P::K1 + vector_index); + float packed[5]; + float value[3]; +#pragma unroll + for (int component = 0; component < 5; ++component) { + packed[component] = probes[3 * P::K1 + component * P::K2 + lane]; + } +#pragma unroll + for (int component = 0; component < 3; ++component) { + value[component] = probes[component * P::K1 + vector_index]; + } + float product[3]; + matrix_vector(packed_to_stf(packed), value, product); + Matrix3 matrix{}; +#pragma unroll + for (int row = 0; row < 3; ++row) { +#pragma unroll + for (int column = 0; column < 3; ++column) { + matrix.value[row][column] = + 2.0f * gradient * product[row] * value[column]; + } + } + float packed_gradient[5]; + matrix_gradient_to_packed(matrix, packed_gradient); +#pragma unroll + for (int component = 0; component < 5; ++component) { + d_tensor[component] += packed_gradient[component]; + } + } + if constexpr (Lmax >= 3) { + accumulate_coupling_gradient(node, args, probes, moments, + 2, lane, scratch); +#pragma unroll + for (int component = 0; component < 5; ++component) { + d_tensor[component] += scratch[component]; + scratch[component] = 0.0f; + } + } +#pragma unroll + for (int component = 0; component < 5; ++component) { + d_probes[3 * P::K1 + component * P::K2 + lane] = d_tensor[component]; + } + } + __syncthreads(); + + __shared__ float d_aligned_storage[P::NodeGroups * P::AlignedWidth]; + float* d_aligned = d_aligned_storage + group * P::AlignedWidth; + for (int channel = lane; channel < P::C1; channel += P::NodeWidth) { + float gradient[3] = {}; + for (int other = 0; other < P::C1; ++other) { + const float upstream = + (channel == other ? 2.0f : kSqrtTwo) * + load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, + P::OutputWidth, + P::OutputGram1 + gram_pair_position(channel, other, P::C1)); +#pragma unroll + for (int component = 0; component < 3; ++component) { + gradient[component] = fmaf(upstream, aligned[component * P::C1 + other], + gradient[component]); + } + } + for (int probe = 0; probe < P::K1; ++probe) { + const float weight = readout_weight(args.readout_matrices, + 5, probe, channel); +#pragma unroll + for (int component = 0; component < 3; ++component) { + gradient[component] = fmaf(weight, d_probes[component * P::K1 + probe], + gradient[component]); + } + } +#pragma unroll + for (int component = 0; component < 3; ++component) { + d_aligned[component * P::C1 + channel] = gradient[component]; + } + } + if (lane < P::C2) { + float gradient[5] = {}; + for (int other = 0; other < P::C2; ++other) { + const float upstream = + (lane == other ? 2.0f : kSqrtTwo) * + load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, + P::OutputWidth, + P::OutputGram2 + gram_pair_position(lane, other, P::C2)); +#pragma unroll + for (int component = 0; component < 5; ++component) { + gradient[component] = + fmaf(upstream, aligned[3 * P::C1 + component * P::C2 + other], + gradient[component]); + } + } + for (int probe = 0; probe < P::K2; ++probe) { + const float weight = + readout_weight(args.readout_matrices, 7, probe, lane); +#pragma unroll + for (int component = 0; component < 5; ++component) { + gradient[component] = + fmaf(weight, d_probes[3 * P::K1 + component * P::K2 + probe], + gradient[component]); + } + } +#pragma unroll + for (int component = 0; component < 5; ++component) { + d_aligned[3 * P::C1 + component * P::C2 + lane] = gradient[component]; + } + } + __syncthreads(); + + __shared__ float d_moments_storage[P::NodeGroups * P::MomentWidth]; + float* d_moments = d_moments_storage + group * P::MomentWidth; + for (int channel = lane; channel < Channels; channel += P::NodeWidth) { + d_moments[P::ScalarOffset + channel] = + load_output_gradient(args.descriptor_gradient, args.output_inv_std, + node, P::OutputWidth, P::OutputScalar + channel); + } + for (int channel = lane; channel < P::C1; channel += P::NodeWidth) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + float value = 0.0f; + for (int output = 0; output < P::C1; ++output) { + value = fmaf(d_aligned[component * P::C1 + output], + readout_weight(args.readout_matrices, 1, + output, channel), + value); + } + d_moments[P::VectorOffset + component * P::C1 + channel] = value; + } + } + if (lane < P::C2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + float value = 0.0f; + for (int output = 0; output < P::C2; ++output) { + value = fmaf(d_aligned[3 * P::C1 + component * P::C2 + output], + readout_weight(args.readout_matrices, 3, + output, lane), + value); + } + d_moments[P::TensorOffset + component * P::C2 + lane] = value; + } + } + if constexpr (Lmax >= 3) { + if (lane < Lmax - 2) { + const int degree = 3 + lane; + const int offset = P::HighOffset + (lane == 0 ? 0 : P::High3); + const int count = 2 * degree + 1; + const float gram = + 2.0f * load_output_gradient(args.descriptor_gradient, + args.output_inv_std, node, P::OutputWidth, + P::OutputGram3 + lane); + accumulate_coupling_gradient(node, args, probes, moments, + degree, 0, scratch); + for (int component = 0; component < count; ++component) { + d_moments[offset + component] = + fmaf(gram, moments[offset + component], scratch[component]); + } + } + } + __syncthreads(); + + // === Normalizer VJPs === + // The scalar and non-scalar blocks carry independent envelope masses, so + // each contributes its own smooth normalizer cotangent. + float scalar_dot = 0.0f; + float angular_dot = 0.0f; + for (int coordinate = lane; coordinate < P::MomentWidth; + coordinate += P::NodeWidth) { + const float product = d_moments[coordinate] * moments[coordinate]; + if (coordinate < Channels) { + scalar_dot += product; + } else { + angular_dot += product; + } + } + const unsigned node_mask = subwarp_mask(group * P::NodeWidth); + scalar_dot = subwarp_sum(scalar_dot, node_mask); + angular_dot = subwarp_sum(angular_dot, node_mask); + // Each mass reaches the output twice: through the moments it normalizes, + // whose cotangent carries the factor -n^2/2, and through its own divisor + // sqrt(mass + floor), whose derivative is n/2. + const float scalar_mass_gradient = + 0.5f * scalar_norm * + (load_output_gradient(args.descriptor_gradient, args.output_inv_std, node, + P::OutputWidth, P::OutputDivisor + 0) - + scalar_dot * scalar_norm); + const float angular_mass_gradient = + 0.5f * angular_norm * + (load_output_gradient(args.descriptor_gradient, args.output_inv_std, node, + P::OutputWidth, P::OutputDivisor + 1) - + angular_dot * angular_norm); + for (int coordinate = lane; coordinate < P::MomentWidth; + coordinate += P::NodeWidth) { + const float value = d_moments[coordinate] * + (coordinate < Channels ? scalar_norm : angular_norm); + if (active) { + args.moment_gradient[node * P::StateWidth + coordinate] = value; + } + } + if (active && lane == 0) { + args.moment_gradient[node * P::StateWidth + P::MomentWidth] = + scalar_mass_gradient; + args.moment_gradient[node * P::StateWidth + P::MomentWidth + 1] = + angular_mass_gradient; + } +} + +// === Edge recomputation backward === + +template +__global__ __launch_bounds__(Profile::Threads, + 32) void edge_backward_kernel(Arguments args) { + using P = Profile; + constexpr int EdgeWidth = P::BackwardEdgeWidth; + constexpr int Groups = kWarpSize / EdgeWidth; + constexpr int ChannelTiles = Channels / EdgeWidth; + constexpr int AngularTiles = (P::C1 + EdgeWidth - 1) / EdgeWidth; + constexpr int TensorTiles = (P::C2 + EdgeWidth - 1) / EdgeWidth; + constexpr int HighTiles = (P::HighCount + EdgeWidth - 1) / EdgeWidth; + + const int thread = threadIdx.x; + const long node = blockIdx.x; + if (node >= args.node_count) { + return; + } + const int base_channel = thread & (EdgeWidth - 1); + const int group = thread / EdgeWidth; + const int leader = group * EdgeWidth; + const unsigned mask = subwarp_mask(leader); + const auto* edge_index = static_cast(args.edge_index); + const auto* destination_order = + static_cast(args.destination_order); + const int center_type = static_cast(args.atype[args.node_begin + node]); + const long begin = args.destination_row_ptr[node]; + const long end = args.destination_row_ptr[node + 1]; + const int radial_modes = HasModes ? args.radial_modes : 0; + const long gradient_offset = node * P::StateWidth; + const float scalar_mass_gradient = + __ldg(args.moment_gradient + gradient_offset + P::MomentWidth); + const float angular_mass_gradient = + __ldg(args.moment_gradient + gradient_offset + P::MomentWidth + 1); + + // The scalar cotangent is one vector of width C0 that every edge rereads. + // Holding it in registers costs one entry per channel tile and spills the + // widest profiles; shared memory serves it as a conflict-free broadcast + // because concurrent edge groups address identical channels. + __shared__ float scalar_gradient[Channels]; + for (int channel = thread; channel < Channels; channel += P::Threads) { + scalar_gradient[channel] = __ldg(args.moment_gradient + gradient_offset + + P::ScalarOffset + channel); + } + float d_vector[AngularTiles][3] = {}; + float d_tensor[TensorTiles][5] = {}; + float d_high[HighTiles > 0 ? HighTiles : 1] = {}; +#pragma unroll + for (int tile = 0; tile < AngularTiles; ++tile) { + const int channel = base_channel + tile * EdgeWidth; + if (channel < P::C1) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + d_vector[tile][component] = + __ldg(args.moment_gradient + gradient_offset + P::VectorOffset + + component * P::C1 + channel); + } + } + } +#pragma unroll + for (int tile = 0; tile < TensorTiles; ++tile) { + const int channel = base_channel + tile * EdgeWidth; + if (channel < P::C2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + d_tensor[tile][component] = + __ldg(args.moment_gradient + gradient_offset + P::TensorOffset + + component * P::C2 + channel); + } + } + } + if constexpr (Lmax >= 3) { +#pragma unroll + for (int component = 0; component < P::HighCount; ++component) { + if (component % EdgeWidth == base_channel) { + d_high[component / EdgeWidth] = __ldg( + args.moment_gradient + gradient_offset + P::HighOffset + component); + } + } + } + + __shared__ float mode_cache[HasModes ? Groups * kModeStride : 1]; + __shared__ float mode_derivative_cache[HasModes ? Groups * kModeStride : 1]; + const int mode_offset = HasModes ? group * kModeStride : 0; + float* modes = mode_cache + mode_offset; + float* mode_derivatives = mode_derivative_cache + mode_offset; + __syncthreads(); + + for (long position = begin + group; position < end; position += Groups) { + const long edge = edge_at_position(position, destination_order); + if (args.edge_mask != nullptr && !args.edge_mask[edge]) { + if (thread == leader) { + args.edge_gradient[edge * 3 + 0] = 0.0f; + args.edge_gradient[edge * 3 + 1] = 0.0f; + args.edge_gradient[edge * 3 + 2] = 0.0f; + } + continue; + } + const EdgeGeometry geometry = load_geometry( + edge, args.rcut, args.eps, args.edge_vec, edge_index, args.atype); + const TableLocation location = + locate_table(geometry.radius, args.table_stride, + args.table_max, args.interval_count); + if constexpr (!Canonical) { + if (center_type >= args.type_count - 1 || + geometry.source_type >= args.type_count - 1) { + if (thread == leader) { + args.edge_gradient[edge * 3 + 0] = 0.0f; + args.edge_gradient[edge * 3 + 1] = 0.0f; + args.edge_gradient[edge * 3 + 2] = 0.0f; + } + continue; + } + } + const TableRow row = table_row(args.table, location, args.table_width); + const float coordinate = location.coordinate; + const bool clamped = location.clamped; + if constexpr (HasModes) { + __syncwarp(mask); + for (int mode = base_channel; mode < radial_modes; mode += EdgeWidth) { + const float2 value = evaluate_table_with_derivative( + row, Channels + mode, coordinate, clamped); + modes[mode] = value.x; + mode_derivatives[mode] = value.y; + } + __syncwarp(mask); + } + + float basis[9]; + fill_angular_basis(geometry, basis); + const float envelope = geometry.envelope; + const long pair = + static_cast(center_type) * args.type_count + geometry.source_type; + const float2* film_row = + reinterpret_cast(args.pair_film + pair * Channels * 2) + + base_channel; + const float* mixing_row = + HasModes + ? args.pair_mixing + (pair * Channels + base_channel) * radial_modes + : nullptr; + + // The high-degree cotangent enters the leading channel only, so its + // angular contraction is reduced onto the lane that owns that channel. + float high_angular = 0.0f; + if constexpr (Lmax >= 3) { +#pragma unroll + for (int component = 0; component < P::HighCount; ++component) { + if (component % EdgeWidth == base_channel) { + high_angular = + fmaf(d_high[component / EdgeWidth], + high_basis_value(geometry, component), high_angular); + } + } + high_angular = subwarp_sum(high_angular, mask); + } + + float radial_gradient = 0.0f; + float envelope_gradient = 0.0f; + float d_basis[9] = {}; + float angular_zero = 0.0f; +#pragma unroll + for (int tile = 0; tile < ChannelTiles; ++tile) { + const int channel = base_channel + tile * EdgeWidth; + const float2 radial = + evaluate_table_with_derivative(row, channel, coordinate, clamped); + const float2 film = __ldg(film_row + tile * EdgeWidth); + // ``film_value`` is the pre-envelope FiLM amplitude phi; the reduced + // payload is phi * chi for degree zero and phi * chi^2 above it. + float film_value = fmaf(film.x, radial.x, film.y); + float film_derivative = film.x * radial.y; + if constexpr (HasModes) { + accumulate_modes_with_derivative( + mixing_row + tile * (EdgeWidth * radial_modes), modes, + mode_derivatives, radial_modes, film_value, film_derivative); + } + float angular = 0.0f; + if (channel < P::C1) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + angular = + fmaf(d_vector[tile][component], basis[1 + component], angular); + } + } + if (channel < P::C2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + angular = + fmaf(d_tensor[tile][component], basis[4 + component], angular); + } + } + if constexpr (Lmax >= 3) { + if (channel == 0) { + angular += high_angular; + } + } + // d/dphi = chi * (p + chi * a) + // d/dchi = phi * (p + 2 * chi * a) + const float scaled = envelope * angular; + const float film_gradient = scalar_gradient[channel] + scaled; + radial_gradient = + fmaf(envelope * film_gradient, film_derivative, radial_gradient); + envelope_gradient = + fmaf(film_gradient + scaled, film_value, envelope_gradient); + const float angular_payload = film_value * envelope * envelope; + if (channel < P::C1) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + d_basis[1 + component] = + fmaf(d_vector[tile][component], angular_payload, + d_basis[1 + component]); + } + } + if (channel < P::C2) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + d_basis[4 + component] = + fmaf(d_tensor[tile][component], angular_payload, + d_basis[4 + component]); + } + } + if constexpr (Lmax >= 3) { + if (channel == 0) { + angular_zero = angular_payload; + } + } + } + + float high_du[3] = {0.0f, 0.0f, 0.0f}; + if constexpr (Lmax >= 3) { + const float amplitude = __shfl_sync(mask, angular_zero, leader); +#pragma unroll + for (int component = 0; component < P::HighCount; ++component) { + if (component % EdgeWidth == base_channel) { + high_basis_gradient(geometry, component, + d_high[component / EdgeWidth] * amplitude, + high_du); + } + } + } + + // Both envelope masses contribute once per edge, so their cotangent joins + // the single lane that also owns the leading channel. + if (thread == leader) { + const float squared = envelope * envelope; + envelope_gradient = + fmaf(4.0f * squared * envelope, angular_mass_gradient, + fmaf(2.0f * envelope, scalar_mass_gradient, envelope_gradient)); + } + radial_gradient = fmaf(envelope_gradient, + c3_envelope_derivative(geometry.radius, args.rcut), + radial_gradient); + + // The basis VJP is linear in the radial and angular cotangents, so + // applying it per lane reduces three Cartesian components instead of the + // full set of angular components across the edge group. + float output[3]; + basis_vjp(geometry, d_basis, high_du, radial_gradient, output); +#pragma unroll + for (int component = 0; component < 3; ++component) { + output[component] = subwarp_sum(output[component], mask); + } + if (thread == leader) { + args.edge_gradient[edge * 3 + 0] = output[0]; + args.edge_gradient[edge * 3 + 1] = output[1]; + args.edge_gradient[edge * 3 + 2] = output[2]; + } + } +} + +template +__global__ void zero_padding_kernel(long node_count, + long edge_count, + const index_t* destination_order, + const long* destination_row_ptr, + float* edge_gradient) { + const long valid_edge_count = destination_row_ptr[node_count]; + for (long position = valid_edge_count + blockIdx.x * blockDim.x + threadIdx.x; + position < edge_count; + position += static_cast(blockDim.x) * gridDim.x) { + const long edge = edge_at_position(position, destination_order); + edge_gradient[edge * 3 + 0] = 0.0f; + edge_gradient[edge * 3 + 1] = 0.0f; + edge_gradient[edge * 3 + 2] = 0.0f; + } +} + +// === Launch dispatch === + +template +struct ForwardLauncher { + static void run(const Arguments& args, cudaStream_t stream) { + using P = Profile; + forward_kernel + <<(args.node_count), P::Threads, 0, stream>>>(args); + } +}; + +template +struct BackwardLauncher { + static void run(const Arguments& args, cudaStream_t stream) { + using P = Profile; + const int node_blocks = + static_cast((args.node_count + P::NodeGroups - 1) / P::NodeGroups); + node_backward_kernel + <<>>(args); + edge_backward_kernel + <<(args.node_count), P::Threads, 0, stream>>>(args); + // The reserved edge slots beyond the physical count are only known on the + // device, so the grid is sized from the storage bound and the surplus + // blocks retire immediately. + if (args.clear_padding) { + constexpr int kPaddingThreads = 128; + constexpr long kPaddingBlockLimit = 1024; + const long padding_blocks = min( + kPaddingBlockLimit, + max(1L, (args.edge_count + kPaddingThreads - 1) / kPaddingThreads)); + zero_padding_kernel + <<(padding_blocks), kPaddingThreads, 0, stream>>>( + args.node_count, args.edge_count, + static_cast(args.destination_order), + args.destination_row_ptr, args.edge_gradient); + } + } +}; + +// A signed and an unsigned 32-bit index share their representation over the +// non-negative range that node and edge identifiers occupy, so the topology +// collapses to two element widths. +template class L> +void dispatch_topology(const Arguments& args, cudaStream_t stream) { + const bool wide = args.index_kind == IndexKind::Bits64; + if (args.canonical) { + if (wide) { + L::run(args, stream); + } else { + L::run(args, stream); + } + } else { + if (wide) { + L::run(args, stream); + } else { + L::run(args, stream); + } + } +} + +// The mode residual is a compile-time specialization for the same reason as +// the angular degree: a descriptor without radial modes must not carry the +// vector temporaries and the shared profile cache of one that has them. +template class L> +void dispatch_modes(const Arguments& args, cudaStream_t stream) { + if (args.radial_modes > 0) { + dispatch_topology(args, stream); + } else { + dispatch_topology(args, stream); + } +} + +// The angular degree is a compile-time specialization because degrees three +// and four add moment accumulators that must not enter the register budget of +// the production ``lmax=2`` path. +// +// The operator entry point validates the degree before it reaches this +// dispatch. The unreachable default is still checked rather than folded into +// the highest degree, so that a degree outside the compiled set can only ever +// fail loudly instead of running a kernel for a different model. +template class L> +void dispatch_degree(const Arguments& args, cudaStream_t stream) { + switch (args.lmax) { + case 2: + dispatch_modes(args, stream); + return; + case 3: + dispatch_modes(args, stream); + return; + case 4: + dispatch_modes(args, stream); + return; + default: + throw std::runtime_error("dpa4c: uncompiled angular degree"); + } +} + +#define DPA4C_DEFINE_CHANNEL(width) \ + void launch_forward_c##width(const Arguments& arguments, \ + cudaStream_t stream) { \ + dispatch_degree(arguments, stream); \ + } \ + void launch_backward_c##width(const Arguments& arguments, \ + cudaStream_t stream) { \ + dispatch_degree(arguments, stream); \ + } + +} // namespace deepmd_dpa4c diff --git a/source/op/pt/dpa4c_graph_compress_launch.h b/source/op/pt/dpa4c_graph_compress_launch.h new file mode 100644 index 0000000000..0cb003dad9 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress_launch.h @@ -0,0 +1,274 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Host-side launch interface of the compressed DPA4C descriptor. +// +// The kernel templates are instantiated once per scalar width in a dedicated +// translation unit, which keeps the compile time of the angular-degree and +// topology specializations bounded and parallel. The dispatching translation +// unit only sees the plain argument bundle declared here. + +#pragma once + +#include + +#include + +namespace deepmd_dpa4c { + +constexpr int kWarpSize = 32; + +/// Element type of the topology indices, folded to two widths because a +/// signed and an unsigned 32-bit index share their representation over the +/// non-negative range that node and edge identifiers occupy. +enum class IndexKind : int { Bits32 = 0, Bits64 = 1 }; + +/// Complete argument bundle of one descriptor launch. +/// +/// Pointers that a given direction does not read stay null. All tensors are +/// contiguous and, except for the topology and the coupling layout, fp32. +struct Arguments { + long node_count = 0; + long edge_count = 0; + int lmax = 2; + int interval_count = 0; + int type_count = 0; + int table_width = 0; + int radial_modes = 0; + // First node of the launched run within the system. Node-indexed buffers are + // already offset by the caller; the atom type table is not, because neighbor + // lookups index it with absolute source indices. + long node_begin = 0; + // Whether this run owns the reserved edge slots past the physical count. + // They lie beyond every destination row, so only the run that ends at the + // last node may clear them; an earlier run would erase the gradients that + // later runs are about to write. + bool clear_padding = true; + int coupling_count = 0; + float table_stride = 0.0f; + float table_max = 0.0f; + float rcut = 0.0f; + float eps = 0.0f; + float degree_floor = 0.0f; + bool canonical = false; + IndexKind index_kind = IndexKind::Bits64; + + const void* edge_index = nullptr; + const void* destination_order = nullptr; + const float* edge_vec = nullptr; + const bool* edge_mask = nullptr; + const long* destination_row_ptr = nullptr; + const long* atype = nullptr; + const float* table = nullptr; + const float* pair_film = nullptr; + const float* pair_mixing = nullptr; + const float* type_embedding = nullptr; + const float* readout_matrices = nullptr; + const int* coupling_meta = nullptr; + const int* coupling_entry = nullptr; + const float* coupling_value = nullptr; + const float* output_mean = nullptr; + const float* output_inv_std = nullptr; + const float* descriptor_gradient = nullptr; + const float* state = nullptr; + + float* descriptor = nullptr; + float* state_out = nullptr; + float* moment_gradient = nullptr; + float* edge_gradient = nullptr; +}; + +// === Compile-time descriptor profile === + +constexpr int triangular(int width) { return width * (width + 1) / 2; } + +constexpr int degree_one_width(int channels) { + int exponent = 0; + for (int value = channels; value > 1; value >>= 1) { + ++exponent; + } + const int width = 1 << ((exponent + 1) / 2); + return width < 4 ? 4 : width; +} + +constexpr int degree_two_width(int channels) { + const int width = degree_one_width(channels) >> 1; + return width < 4 ? 4 : width; +} + +constexpr int degree_rank(int degree, int rank_one, int rank_two) { + return degree == 1 ? rank_one : (degree == 2 ? rank_two : 1); +} + +// Independent probe contractions emitted by one degree triple. Axes carrying +// equal degrees are symmetric under the symmetrized coupling, so only one +// representative of each orbit is emitted. +constexpr int triple_outputs(int l1, int l2, int l3, int k1, int k2) { + const int r1 = degree_rank(l1, k1, k2); + const int r2 = degree_rank(l2, k1, k2); + const int r3 = degree_rank(l3, k1, k2); + if (l1 == l3) { + return r1 * (r1 + 1) * (r1 + 2) / 6; + } + if (l1 == l2) { + return triangular(r1) * r3; + } + if (l2 == l3) { + return r1 * triangular(r2); + } + return r1 * r2 * r3; +} + +// Outputs emitted by every triple enumerated before the target. Passing a +// triple that cannot occur returns the complete bispectrum width. +constexpr int bispectrum_prefix(int lmax, int k1, int k2, int a, int b, int c) { + int total = 0; + for (int l1 = 1; l1 <= lmax; ++l1) { + for (int l2 = l1; l2 <= lmax; ++l2) { + for (int l3 = l2; l3 <= lmax; ++l3) { + if (l3 > l1 + l2 || (l1 + l2 + l3) % 2 != 0) { + continue; + } + if (l1 == a && l2 == b && l3 == c) { + return total; + } + total += triple_outputs(l1, l2, l3, k1, k2); + } + } + } + return total; +} + +// Degree triples that the sparse coupling artifact must describe, that is +// every allowed triple except the two the kernel contracts in closed form. +constexpr int coupling_record_count(int lmax) { + int total = 0; + for (int l1 = 1; l1 <= lmax; ++l1) { + for (int l2 = l1; l2 <= lmax; ++l2) { + for (int l3 = l2; l3 <= lmax; ++l3) { + if (l3 > l1 + l2 || (l1 + l2 + l3) % 2 != 0) { + continue; + } + if ((l1 == 1 && l2 == 1 && l3 == 2) || + (l1 == 2 && l2 == 2 && l3 == 2)) { + continue; + } + ++total; + } + } + } + return total; +} + +// Number of lanes that cooperate on one edge. A warp therefore keeps +// ``32 / width`` edges in flight, and the per-edge geometry, which every lane +// of a group recomputes, is amortized over that many edges. Narrowing the +// group lowers that fixed cost but raises the per-lane channel and moment +// footprint, because an angular channel beyond the group width has to be +// tiled and consumes another accumulator. Each profile therefore sits at the +// measured optimum of that trade-off on a diamond neighborhood; forward and +// backward differ because the backward carries the additional angular +// cotangents. +template +struct EdgeMap; + +template <> +struct EdgeMap<8> { + static constexpr int Forward = 2; + static constexpr int Backward = 2; +}; +template <> +struct EdgeMap<16> { + static constexpr int Forward = 4; + static constexpr int Backward = 4; +}; +template <> +struct EdgeMap<32> { + static constexpr int Forward = 8; + static constexpr int Backward = 4; +}; +template <> +struct EdgeMap<64> { + static constexpr int Forward = 8; + static constexpr int Backward = 8; +}; +template <> +struct EdgeMap<128> { + static constexpr int Forward = 16; + static constexpr int Backward = 8; +}; + +template +struct Profile { + static constexpr int C0 = Channels; + static constexpr int C1 = degree_one_width(Channels); + static constexpr int C2 = degree_two_width(Channels); + static constexpr int K1 = C2; + static constexpr int K2 = 2; + + // Flat moment layout: degree zero, degree one, degree two, then the + // single-channel high degrees in increasing order. + static constexpr int ScalarOffset = 0; + static constexpr int VectorOffset = C0; + static constexpr int TensorOffset = C0 + 3 * C1; + static constexpr int HighOffset = TensorOffset + 5 * C2; + static constexpr int High3 = Lmax >= 3 ? 7 : 0; + static constexpr int High4 = Lmax >= 4 ? 9 : 0; + static constexpr int HighCount = High3 + High4; + static constexpr int MomentWidth = HighOffset + HighCount; + static constexpr int StateWidth = MomentWidth + 2; + + // Cached intermediates of the invariant readout. + static constexpr int AlignedWidth = 3 * C1 + 5 * C2; + static constexpr int ProbeWidth = 3 * K1 + 5 * K2; + + // Descriptor layout. + static constexpr int Gram1 = triangular(C1); + static constexpr int Gram2 = triangular(C2); + static constexpr int Bis112 = triangular(K1) * K2; + static constexpr int Bis222 = 4; + static constexpr int Quartic = K1 * K2; + static constexpr int OutputScalar = 0; + static constexpr int OutputGram1 = C0; + static constexpr int OutputGram2 = OutputGram1 + Gram1; + static constexpr int OutputGram3 = OutputGram2 + Gram2; + static constexpr int OutputGram4 = OutputGram3 + (Lmax >= 3 ? 1 : 0); + static constexpr int BispectrumBase = OutputGram4 + (Lmax >= 4 ? 1 : 0); + static constexpr int OutputBis112 = + BispectrumBase + bispectrum_prefix(Lmax, K1, K2, 1, 1, 2); + static constexpr int OutputBis222 = + BispectrumBase + bispectrum_prefix(Lmax, K1, K2, 2, 2, 2); + static constexpr int OutputQuartic = + BispectrumBase + bispectrum_prefix(Lmax, K1, K2, 0, 0, 0); + // The two moment divisors close the geometric block. Normalization is + // otherwise irreversible, so without them neither the readout nor the + // fitting network can see the effective coordination they encode. + static constexpr int OutputDivisor = OutputQuartic + Quartic; + static constexpr int OutputType = OutputDivisor + 2; + static constexpr int OutputWidth = OutputType + C0; + + // The group width does not widen for the high angular degrees. Their + // single-channel components add one accumulator per lane and tile, but every + // measured widening lost more to the reduced edge concurrency than it + // recovered in register pressure. + static constexpr int ForwardEdgeWidth = EdgeMap::Forward; + static constexpr int BackwardEdgeWidth = EdgeMap::Backward; + static constexpr int NodeWidth = 8; + static constexpr int NodeGroups = kWarpSize / NodeWidth; + static constexpr int Threads = kWarpSize; +}; + +/// Scalar widths that own a compiled specialization. +#define DPA4C_FOR_EACH_CHANNEL(macro) \ + macro(8) macro(16) macro(32) macro(64) macro(128) + +#define DPA4C_DECLARE_CHANNEL(width) \ + void launch_forward_c##width(const Arguments& arguments, \ + cudaStream_t stream); \ + void launch_backward_c##width(const Arguments& arguments, \ + cudaStream_t stream); + +DPA4C_FOR_EACH_CHANNEL(DPA4C_DECLARE_CHANNEL) + +#undef DPA4C_DECLARE_CHANNEL + +} // namespace deepmd_dpa4c diff --git a/source/op/pt/edge_force_virial.cu b/source/op/pt/edge_force_virial.cu index af0eb67861..5dff4b4990 100644 --- a/source/op/pt/edge_force_virial.cu +++ b/source/op/pt/edge_force_virial.cu @@ -18,6 +18,7 @@ #include #include +#include #include namespace { @@ -146,27 +147,33 @@ __global__ void edge_force_virial_kernel( } } -template -__global__ void reduce_node_virial_kernel( +// Partial pass of a segment sum over the node axis. Each frame owns a +// contiguous span of that axis, so a block can reduce a strided slice of one +// (frame, component) pair without any atomic, and the reduction order follows +// the launch geometry alone and is therefore reproducible. Accumulation is +// fp64 whatever the stored type. +template +__global__ void reduce_frame_segment_kernel( long frame_count, int partial_count, const long* __restrict__ frame_row_ptr, - const scalar_t* __restrict__ node_virial, + const scalar_t* __restrict__ node_values, double* __restrict__ partial) { __shared__ double values[kThreads]; - const long task_count = static_cast(frame_count) * 9 * partial_count; + const long task_count = + static_cast(frame_count) * kComponents * partial_count; for (long task = blockIdx.x; task < task_count; task += gridDim.x) { const int partial_index = task % partial_count; const long output = task / partial_count; - const long frame = output / 9; - const int component = output % 9; + const long frame = output / kComponents; + const int component = output % kComponents; const long begin = frame_row_ptr[frame]; const long end = frame_row_ptr[frame + 1]; double sum = 0.0; for (long node = begin + partial_index * static_cast(blockDim.x) + threadIdx.x; node < end; node += static_cast(partial_count) * blockDim.x) { - sum += static_cast(node_virial[node * 9 + component]); + sum += static_cast(node_values[node * kComponents + component]); } values[threadIdx.x] = sum; __syncthreads(); @@ -183,11 +190,13 @@ __global__ void reduce_node_virial_kernel( } } +// Final pass of a segment sum: fold the per-block partials of each output. template -__global__ void finalize_virial_kernel(long output_count, - int partial_count, - const double* __restrict__ partial, - scalar_t* __restrict__ virial) { +__global__ void finalize_frame_segment_kernel( + long output_count, + int partial_count, + const double* __restrict__ partial, + scalar_t* __restrict__ out) { __shared__ double values[kThreads]; for (long output = blockIdx.x; output < output_count; output += gridDim.x) { double sum = 0.0; @@ -203,16 +212,61 @@ __global__ void finalize_virial_kernel(long output_count, __syncthreads(); } if (threadIdx.x == 0) { - virial[output] = static_cast(values[0]); + out[output] = static_cast(values[0]); } __syncthreads(); } } +// Exclusive prefix of the per-frame node counts, giving each frame's span of +// the node axis as ``[row_ptr[f], row_ptr[f + 1])``. +torch::Tensor frame_row_pointer(const torch::Tensor& n_node_per_frame) { + return torch::cat({torch::zeros({1}, n_node_per_frame.options()), + torch::cumsum(n_node_per_frame, 0)}) + .to(torch::kInt64) + .contiguous(); +} + +// Slices per frame that the partial pass reduces in parallel. Bounding it +// keeps the partial buffer small for a single large frame, which is the +// molecular dynamics case. +int frame_segment_partials(long node_count, long frame_count) { + const long average_node_count = (node_count + frame_count - 1) / frame_count; + return static_cast( + std::min((average_node_count + kThreads - 1) / kThreads, + static_cast(kMaximumVirialPartials))); +} + +// Sum ``node_values``, laid out as (node, component), over the node segment of +// each frame. Returns a flat (frame * kComponents) buffer. +template +void launch_frame_segment_sum(long node_count, + long frame_count, + const torch::Tensor& frame_row_ptr, + const scalar_t* node_values, + torch::Tensor& partial, + scalar_t* out, + cudaStream_t stream) { + const int partial_count = static_cast(partial.size(1)); + const long output_count = frame_count * kComponents; + + const int partial_blocks = std::min(output_count * partial_count, 65535L); + reduce_frame_segment_kernel + <<>>( + frame_count, partial_count, frame_row_ptr.data_ptr(), + node_values, partial.data_ptr()); + FORCE_CHECK_LAUNCH("frame segment sum partial reduction"); + + const int final_blocks = std::min(output_count, 65535L); + finalize_frame_segment_kernel + <<>>(output_count, partial_count, + partial.data_ptr(), out); + FORCE_CHECK_LAUNCH("frame segment sum final reduction"); +} + template void launch_force_virial(long node_count, long frame_count, - int partial_count, const torch::Tensor& edge_gradient, const torch::Tensor& edge_vec, const torch::Tensor& edge_mask, @@ -240,18 +294,9 @@ void launch_force_virial(long node_count, force.data_ptr(), node_virial.data_ptr()); FORCE_CHECK_LAUNCH("edge_force_virial node reduction"); - const long output_count = static_cast(frame_count) * 9; - const int partial_blocks = std::min(output_count * partial_count, 65535L); - reduce_node_virial_kernel<<>>( - frame_count, partial_count, frame_row_ptr.data_ptr(), - node_virial.data_ptr(), virial_partial.data_ptr()); - FORCE_CHECK_LAUNCH("edge_force_virial partial frame reduction"); - - const int final_blocks = std::min(output_count, 65535L); - finalize_virial_kernel<<>>( - output_count, partial_count, virial_partial.data_ptr(), - virial.data_ptr()); - FORCE_CHECK_LAUNCH("edge_force_virial final frame reduction"); + launch_frame_segment_sum( + node_count, frame_count, frame_row_ptr, node_virial.data_ptr(), + virial_partial, virial.data_ptr(), stream); } std::tuple assemble_force_virial( @@ -278,15 +323,8 @@ std::tuple assemble_force_virial( return {force, atom_virial, virial}; } - auto frame_row_ptr = - torch::cat({torch::zeros({1}, n_node_per_frame.options()), - torch::cumsum(n_node_per_frame, 0)}) - .to(torch::kInt64) - .contiguous(); - const long average_node_count = (node_count + frame_count - 1) / frame_count; - const int partial_count = - static_cast(std::min((average_node_count + kThreads - 1) / kThreads, - static_cast(kMaximumVirialPartials))); + auto frame_row_ptr = frame_row_pointer(n_node_per_frame); + const int partial_count = frame_segment_partials(node_count, frame_count); auto virial_partial = torch::empty({frame_count * 9, partial_count}, options.dtype(torch::kFloat64)); const auto stream = at::cuda::getCurrentCUDAStream(); @@ -295,14 +333,20 @@ std::tuple assemble_force_virial( edge_gradient.scalar_type(), "edge_force_virial", [&] { if (source_order.scalar_type() == torch::kInt32) { launch_force_virial( - node_count, frame_count, partial_count, edge_gradient, edge_vec, - edge_mask, destination_order, destination_row_ptr, source_order, + node_count, frame_count, edge_gradient, edge_vec, edge_mask, + destination_order, destination_row_ptr, source_order, + source_row_ptr, frame_row_ptr, force, node_virial, virial_partial, + virial, stream); + } else if (source_order.scalar_type() == torch::kUInt32) { + launch_force_virial( + node_count, frame_count, edge_gradient, edge_vec, edge_mask, + destination_order, destination_row_ptr, source_order, source_row_ptr, frame_row_ptr, force, node_virial, virial_partial, virial, stream); } else { launch_force_virial( - node_count, frame_count, partial_count, edge_gradient, edge_vec, - edge_mask, destination_order, destination_row_ptr, source_order, + node_count, frame_count, edge_gradient, edge_vec, edge_mask, + destination_order, destination_row_ptr, source_order, source_row_ptr, frame_row_ptr, force, node_virial, virial_partial, virial, stream); } @@ -393,10 +437,11 @@ std::tuple edge_force_virial( "edge_force_virial: CSR tensors must be contiguous"); TORCH_CHECK( (source_order.scalar_type() == torch::kInt32 || + source_order.scalar_type() == torch::kUInt32 || source_order.scalar_type() == torch::kInt64) && destination_order.scalar_type() == source_order.scalar_type(), "edge_force_virial: destination_order and source_order must have the " - "same int32 or int64 dtype"); + "same int32, uint32, or int64 dtype"); return assemble_force_virial(node_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, n_node_per_frame, @@ -430,9 +475,10 @@ canonical_edge_force_virial(torch::Tensor edge_gradient, source_row_ptr.scalar_type() == torch::kInt64, "canonical_edge_force_virial: row pointers must be int64"); TORCH_CHECK(source_order.scalar_type() == torch::kInt32 || + source_order.scalar_type() == torch::kUInt32 || source_order.scalar_type() == torch::kInt64, - "canonical_edge_force_virial: source_order must be int32 or " - "int64"); + "canonical_edge_force_virial: source_order must be int32, " + "uint32, or int64"); TORCH_CHECK(destination_row_ptr.numel() == node_count + 1 && source_row_ptr.numel() == node_count + 1, "canonical_edge_force_virial: row pointers must have N + 1 " @@ -446,6 +492,49 @@ canonical_edge_force_virial(torch::Tensor edge_gradient, want_atom_virial); } +// Per-frame total of a scalar carried on the node axis, the energy being the +// only such quantity. Scattering it with an index add would serialize one +// atomic per node on a single address whenever the batch holds one frame, +// which is the molecular dynamics case; the segment reduction that already +// serves the virial has no atomic and no node-length index to materialize. +// +// The frames cover ``[0, sum(n_node_per_frame))`` of the node axis; nodes past +// that are padding and contribute to no frame. The caller owns that +// invariant, since checking it would force a device synchronization. +torch::Tensor frame_scalar_sum(torch::Tensor node_scalar, + torch::Tensor n_node_per_frame) { + const long node_count = node_scalar.size(0); + const long frame_count = n_node_per_frame.size(0); + TORCH_CHECK(node_scalar.is_cuda() && n_node_per_frame.is_cuda(), + "frame_scalar_sum: inputs must be CUDA tensors"); + TORCH_CHECK(node_scalar.device() == n_node_per_frame.device(), + "frame_scalar_sum: inputs must share one device"); + TORCH_CHECK(node_scalar.is_contiguous(), + "frame_scalar_sum: node_scalar must be contiguous"); + TORCH_CHECK(node_scalar.dim() == 2 && node_scalar.size(1) == 1, + "frame_scalar_sum: node_scalar must have shape (N, 1)"); + TORCH_CHECK(n_node_per_frame.dim() == 1, + "frame_scalar_sum: n_node_per_frame must be one-dimensional"); + + auto out = torch::zeros({frame_count, 1}, node_scalar.options()); + if (node_count == 0 || frame_count == 0) { + return out; + } + auto frame_row_ptr = frame_row_pointer(n_node_per_frame); + const int partial_count = frame_segment_partials(node_count, frame_count); + auto partial = torch::empty({frame_count, partial_count}, + node_scalar.options().dtype(torch::kFloat64)); + const auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(node_scalar.scalar_type(), "frame_scalar_sum", + [&] { + launch_frame_segment_sum( + node_count, frame_count, frame_row_ptr, + node_scalar.data_ptr(), partial, + out.data_ptr(), stream); + }); + return out; +} + TORCH_LIBRARY_FRAGMENT(deepmd, library) { library.def( "build_graph_csr(Tensor edge_index, SymInt node_count, " @@ -469,10 +558,15 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "(Tensor force, Tensor atom_virial, Tensor virial)"); library.impl("canonical_edge_force_virial", torch::kCUDA, &canonical_edge_force_virial); + library.def( + "frame_scalar_sum(Tensor node_scalar, Tensor n_node_per_frame) " + "-> Tensor"); + library.impl("frame_scalar_sum", torch::kCUDA, &frame_scalar_sum); } TORCH_LIBRARY_IMPL(deepmd, Autograd, library) { library.impl("edge_force_virial", torch::CppFunction::makeFallthrough()); library.impl("canonical_edge_force_virial", torch::CppFunction::makeFallthrough()); + library.impl("frame_scalar_sum", torch::CppFunction::makeFallthrough()); } diff --git a/source/op/pt/graph_fitting.cu b/source/op/pt/graph_fitting.cu index 60724679c8..a9e38a4d18 100644 --- a/source/op/pt/graph_fitting.cu +++ b/source/op/pt/graph_fitting.cu @@ -3,32 +3,41 @@ // Fused energy fitting network for graph-lower inference. The operator is // descriptor-agnostic: any graph-lowered energy model whose fitting is a // plain MLP over the flat node axis dispatches here. -// h_0 = act(x @ W_0 + b_0) * idt_0 (+ identity residual when -// h_l = act(h_{l-1} @ W_l + b_l) * idt_l the layer is square) +// h_0 = act(x @ W_0 + b_0) (+ identity residual when +// h_l = act(h_{l-1} @ W_l + b_l) the layer is square) // e = h_{L-1} @ w_head + b_head + bias_atom_e[atype] (fp64 output) // The GEMMs run on cuBLAS in pedantic fp32 (TF32 off); each layer's bias, -// activation, timestep and residual collapse into one elementwise epilogue -// kernel that also stores the activation derivative for the backward. The -// backward (upstream d_e, a unit vector for the energy reduction) chains +// activation and residual collapse into one elementwise epilogue kernel. +// The backward (upstream d_e, a unit vector for the energy reduction) chains // dh_{L-1} = d_e * w_head^T -// dpre_l = dh_l * act'_l +// dpre_l = dh_l * act'(pre_l + b_l) // dh_{l-1} = dpre_l @ W_l^T (+ dh_l identity residual) // d_x = dpre_0 @ W_0^T // with the elementwise steps fused likewise. // +// The saved state is the pre-activation of every layer, which each GEMM writes +// directly. The backward re-derives the activation derivative from it, so the +// epilogue writes only the activation and the layer costs one full-tensor +// store less than a formulation that also materializes the derivative. +// // All tensors here are node-scale (atoms, not edges); the fusion removes // kernel launches and aten glue rather than FLOPs. The head bias arrives as // a device tensor so that symbolic tracing never reads a value host-side. #include +#include #include #include #include +#include +#include #include #include #include +#include "graph_ops.h" + namespace { #define FITTING_CHECK_LAUNCH(what) \ @@ -82,46 +91,78 @@ void gemm_nt(cudaStream_t stream, a, k, &beta, c, n); } -// Activation value and derivative; codes follow -// deepmd.kernels.triton.dpa1.activation.ACT_CODES (0 = tanh, 1 = silu). +// sigma(z) = (1 + tanh(z/2)) / 2. The identity is not the detour it looks +// like: CUDA's tanhf is a branch-free sequence of one MUFU.EX2 and one +// MUFU.RCP, whereas an accurate expf carries a branchy slow path and the +// intrinsic __expf still needs a reciprocal to finish the sigmoid. Measured +// per kernel instance at one million nodes, the identity beats both, and the +// only faster form is fully approximate, worth 0.5% of the fitting forward +// against an approximate division and exponential. +__device__ __forceinline__ float sigmoid(float z) { + return 0.5f * (1.f + tanhf(0.5f * z)); +} + +// Activation codes follow deepmd.kernels.triton.dpa1.activation.ACT_CODES +// (0 = tanh, 1 = silu). Value and derivative are separate because the forward +// needs only the former and the backward only the latter. +template +__device__ __forceinline__ float act_value(float z) { + if constexpr (ACT == 0) { + return tanhf(z); + } + return z * sigmoid(z); +} + template -__device__ __forceinline__ float2 act_vg(float z) { +__device__ __forceinline__ float act_derivative(float z) { if constexpr (ACT == 0) { - float a = tanhf(z); - return make_float2(a, 1.f - a * a); + const float a = tanhf(z); + return 1.f - a * a; } - const float s = 0.5f * (1.f + tanhf(0.5f * z)); // sigmoid via tanh identity - return make_float2(z * s, s * (1.f + z * (1.f - s))); + const float s = sigmoid(z); + return s * (1.f + z * (1.f - s)); +} + +long ceil_div(long a, long b) { return (a + b - 1) / b; } + +// The node axis is long, so deriving a channel from a flat element index +// costs a 64-bit division per thread. The elementwise epilogues instead index +// (channel, node) directly: threadIdx.x selects the float4 lane within a row, +// keeping a warp's accesses contiguous, and the remaining dimensions walk the +// node axis. Layer widths are multiples of four (Python gate). +struct ElementwiseLaunch { + dim3 grid; + dim3 block; +}; + +ElementwiseLaunch elementwise_launch(long n_node, int dout) { + const int lanes = dout / 4; + const int rows = std::max(1, 256 / lanes); + return {dim3((unsigned)ceil_div(n_node, rows)), dim3(lanes, rows)}; } -// y = act(pre + b) * idt (+ x residual); adot = act' * idt is stored for the -// backward. float4 lanes; layer widths are multiples of four (Python gate). +// y = act(pre + b) (+ x residual). The pre-activation stays where the GEMM +// wrote it, so the layer stores one tensor instead of two. template -__global__ void layer_epilogue_kernel(long total4, +__global__ void layer_epilogue_kernel(long n_node, int dout, const float* __restrict__ pre, const float* __restrict__ b, - const float* __restrict__ idt, const float* __restrict__ x, int residual, - float* __restrict__ y, - float* __restrict__ adot) { - const long t4 = blockIdx.x * (long)blockDim.x + threadIdx.x; - if (t4 >= total4) { + float* __restrict__ y) { + const long node = blockIdx.x * (long)blockDim.y + threadIdx.y; + if (node >= n_node) { return; } - const long t = t4 * 4; - const int c = (int)(t % dout); + const int c = (int)threadIdx.x * 4; + const long t = node * dout + c; const float4 p = *reinterpret_cast(pre + t); const float4 bb = b ? *reinterpret_cast(b + c) : make_float4(0, 0, 0, 0); - const float4 ii = - idt ? *reinterpret_cast(idt + c) : make_float4(1, 1, 1, 1); - const float2 v0 = act_vg(p.x + bb.x); - const float2 v1 = act_vg(p.y + bb.y); - const float2 v2 = act_vg(p.z + bb.z); - const float2 v3 = act_vg(p.w + bb.w); - float4 yy = make_float4(v0.x * ii.x, v1.x * ii.y, v2.x * ii.z, v3.x * ii.w); + float4 yy = + make_float4(act_value(p.x + bb.x), act_value(p.y + bb.y), + act_value(p.z + bb.z), act_value(p.w + bb.w)); if (residual) { const float4 xx = *reinterpret_cast(x + t); yy.x += xx.x; @@ -130,8 +171,6 @@ __global__ void layer_epilogue_kernel(long total4, yy.w += xx.w; } *reinterpret_cast(y + t) = yy; - *reinterpret_cast(adot + t) = - make_float4(v0.y * ii.x, v1.y * ii.y, v2.y * ii.z, v3.y * ii.w); } // Energy head: e[n] = h[n] @ w_head + b_head + bias_atom_e[atype[n]]. @@ -157,138 +196,241 @@ __global__ void head_kernel(long n_node, e[n] = (double)(acc + (b_head ? b_head[0] : 0.f)) + bias_atom_e[atype[n]]; } -// Backward seed: dh_{L-1}[n, c] = d_e[n] * w_head[c] (fp64 upstream). -__global__ void seed_kernel(long total4, - int dout, - const double* __restrict__ d_e, - const float* __restrict__ w_head, - float* __restrict__ dh) { - const long t4 = blockIdx.x * (long)blockDim.x + threadIdx.x; - if (t4 >= total4) { +// Construct the head gradient and apply the final hidden activation VJP in one +// pass. A square final layer also preserves the unmodified head gradient for +// its identity branch. +template +__global__ void seed_backward_epilogue_kernel( + long n_node, + int dout, + const double* __restrict__ d_e, + const float* __restrict__ w_head, + const float* __restrict__ pre, + const float* __restrict__ b, + float* __restrict__ dpre, + float* __restrict__ residual_out) { + const long node = blockIdx.x * (long)blockDim.y + threadIdx.y; + if (node >= n_node) { return; } - const long t = t4 * 4; - const long n = t / dout; - const int c = (int)(t % dout); - const float de = (float)d_e[n]; + const int c = (int)threadIdx.x * 4; + const long t = node * dout + c; + const float de = (float)d_e[node]; const float4 wv = *reinterpret_cast(w_head + c); - *reinterpret_cast(dh + t) = - make_float4(de * wv.x, de * wv.y, de * wv.z, de * wv.w); + const float4 dh = make_float4(de * wv.x, de * wv.y, de * wv.z, de * wv.w); + const float4 p = *reinterpret_cast(pre + t); + const float4 bb = + b ? *reinterpret_cast(b + c) : make_float4(0, 0, 0, 0); + if (residual_out) { + *reinterpret_cast(residual_out + t) = dh; + } + *reinterpret_cast(dpre + t) = + make_float4(dh.x * act_derivative(p.x + bb.x), + dh.y * act_derivative(p.y + bb.y), + dh.z * act_derivative(p.z + bb.z), + dh.w * act_derivative(p.w + bb.w)); } -// Convert dh to dpre in place: dh *= adot. -__global__ void backward_epilogue_kernel(long total4, +// Convert dh to dpre in place: dh *= act'(pre + b). A square residual layer +// also preserves the unmodified dh in its output buffer before the GEMM +// accumulates the weighted branch with beta = 1. +template +__global__ void backward_epilogue_kernel(long n_node, + int dout, const float* __restrict__ dh, - const float* __restrict__ adot, - float* __restrict__ dpre) { - const long t4 = blockIdx.x * (long)blockDim.x + threadIdx.x; - if (t4 >= total4) { + const float* __restrict__ pre, + const float* __restrict__ b, + float* __restrict__ dpre, + float* __restrict__ residual_out) { + const long node = blockIdx.x * (long)blockDim.y + threadIdx.y; + if (node >= n_node) { return; } - const long t = t4 * 4; + const int c = (int)threadIdx.x * 4; + const long t = node * dout + c; const float4 d = *reinterpret_cast(dh + t); - const float4 a = *reinterpret_cast(adot + t); + const float4 p = *reinterpret_cast(pre + t); + const float4 bb = + b ? *reinterpret_cast(b + c) : make_float4(0, 0, 0, 0); + if (residual_out) { + *reinterpret_cast(residual_out + t) = d; + } *reinterpret_cast(dpre + t) = - make_float4(d.x * a.x, d.y * a.y, d.z * a.z, d.w * a.w); + make_float4(d.x * act_derivative(p.x + bb.x), + d.y * act_derivative(p.y + bb.y), + d.z * act_derivative(p.z + bb.z), + d.w * act_derivative(p.w + bb.w)); } -long ceil_div(long a, long b) { return (a + b - 1) / b; } +/// Dispatch a kernel template over the two supported activation codes. +template +void dispatch_activation(long act, Fn&& launch) { + if (act == 0) { + launch(std::integral_constant{}); + } else { + launch(std::integral_constant{}); + } +} } // namespace +FittingLayerPlan fitting_layer_plan(const std::vector& ws) { + FittingLayerPlan plan{std::vector(ws.size() + 1, 0), 0, (int)ws.size()}; + for (size_t l = 0; l < ws.size(); ++l) { + plan.offset[l + 1] = plan.offset[l] + ws[l].size(1); + plan.width_max = std::max(plan.width_max, (long)ws[l].size(1)); + } + return plan; +} + +// Evaluate the network over one contiguous run of nodes. Every full-width +// tensor is indexed from the run's first node, so the same code serves the +// whole node axis and a single tile of it. ``saved`` and ``activation`` are +// sized for the run, not for the system. +void fitting_forward_range(cudaStream_t stream, + const FittingLayerPlan& plan, + const float* x, + long input_width, + const long* atype, + const std::vector& ws, + const std::vector& bs, + const std::vector& resnets, + const torch::Tensor& w_head, + const torch::Tensor& b_head, + const torch::Tensor& bias_atom_e, + int64_t act, + long run_nodes, + float* saved, + float* const activation[2], + double* e) { + const float* cur = x; + int din = (int)input_width; + for (int l = 0; l < plan.n_layer; ++l) { + const int dout = (int)ws[l].size(1); + float* pre = saved + plan.offset[l] * run_nodes; + float* y = activation[l & 1]; + gemm_nn(stream, cur, ws[l].data_ptr(), pre, (int)run_nodes, dout, + din); + const ElementwiseLaunch shape = elementwise_launch(run_nodes, dout); + const bool residual = resnets[l] && dout == din; + dispatch_activation(act, [&](auto tag) { + layer_epilogue_kernel + <<>>( + run_nodes, dout, pre, + bs[l].numel() ? bs[l].data_ptr() : nullptr, cur, + residual ? 1 : 0, y); + }); + FITTING_CHECK_LAUNCH("graph_fitting layer"); + cur = y; + din = dout; + } + head_kernel<<>>( + run_nodes, din, cur, w_head.data_ptr(), + b_head.numel() ? b_head.data_ptr() : nullptr, + bias_atom_e.data_ptr(), atype, e); + FITTING_CHECK_LAUNCH("graph_fitting head"); +} + +// Propagate the head cotangent of one run of nodes back to the input. ``dh`` +// and ``dh_next`` are scratch of the run's size; ``d_x`` is indexed from the +// run's first node. +void fitting_backward_range(cudaStream_t stream, + const FittingLayerPlan& plan, + const double* d_e, + const float* saved, + const std::vector& ws, + const std::vector& bs, + const std::vector& resnets, + const torch::Tensor& w_head, + int64_t act, + long run_nodes, + float* dh, + float* dh_next, + float* d_x) { + for (int l = plan.n_layer - 1; l >= 0; --l) { + const int dout = (int)ws[l].size(1); + const int din = (int)ws[l].size(0); + const float* pre = saved + plan.offset[l] * run_nodes; + const float* b = bs[l].numel() ? bs[l].data_ptr() : nullptr; + float* out = l > 0 ? dh_next : d_x; + const bool residual = resnets[l] && dout == din; + const ElementwiseLaunch shape = elementwise_launch(run_nodes, dout); + dispatch_activation(act, [&](auto tag) { + constexpr int kAct = decltype(tag)::value; + if (l == plan.n_layer - 1) { + seed_backward_epilogue_kernel + <<>>( + run_nodes, dout, d_e, w_head.data_ptr(), pre, b, dh, + residual ? out : nullptr); + } else { + backward_epilogue_kernel<<>>( + run_nodes, dout, dh, pre, b, dh, residual ? out : nullptr); + } + }); + FITTING_CHECK_LAUNCH("graph_fitting backward layer"); + gemm_nt(stream, dh, ws[l].data_ptr(), out, (int)run_nodes, din, dout, + residual ? 1.f : 0.f); + if (l > 0) { + std::swap(dh, dh_next); + } + } +} + // Forward: per-atom energy (fp64 (N, 1)) plus the flat saved buffer of the -// activation derivatives -- adot chunks, chunk l a contiguous (N, width_l) -// sheet. The backward needs only these derivatives; the activations themselves -// stay in a forward-only ping-pong. +// layer pre-activations -- chunk l a contiguous (N, width_l) sheet. The +// activations themselves stay in a forward-only ping-pong. std::tuple graph_fitting( torch::Tensor x, torch::Tensor atype, std::vector ws, std::vector bs, - std::vector idts, std::vector resnets, torch::Tensor w_head, torch::Tensor b_head, torch::Tensor bias_atom_e, int64_t act) { + TORCH_CHECK(x.is_cuda(), "graph_fitting: x must be a CUDA tensor"); + const c10::cuda::CUDAGuard device_guard(x.device()); const long n_node = x.size(0); - auto stream = at::cuda::getCurrentCUDAStream(); - const int n_layer = (int)ws.size(); - std::vector offset(n_layer + 1, 0); - for (int l = 0; l < n_layer; ++l) { - offset[l + 1] = offset[l] + ws[l].size(1); - } - const long total_width = offset[n_layer]; + const FittingLayerPlan plan = fitting_layer_plan(ws); auto f32 = x.options().dtype(torch::kFloat32); - auto saved = torch::empty({n_node * total_width}, f32); + auto saved = torch::empty({n_node * plan.saved_width()}, f32); auto e = torch::empty({n_node, 1}, x.options().dtype(torch::kFloat64)); if (n_node == 0) { return {e, saved}; } - - long width_max = 0; - for (int l = 0; l < n_layer; ++l) { - width_max = std::max(width_max, (long)ws[l].size(1)); - } - // Two-slot ping-pong for the activations: layer l writes slot ``l & 1`` while - // reading the previous layer's slot, so an activation is overwritten only - // after the next GEMM has consumed it (kernels run in stream order). - auto act_buf = torch::empty({2, n_node, width_max}, f32); - float* act_slot[2] = {act_buf[0].data_ptr(), - act_buf[1].data_ptr()}; - const float* cur = x.data_ptr(); - int din = (int)x.size(1); - for (int l = 0; l < n_layer; ++l) { - const int dout = (int)ws[l].size(1); - float* h = act_slot[l & 1]; - float* adot = saved.data_ptr() + offset[l] * n_node; - gemm_nn(stream, cur, ws[l].data_ptr(), h, (int)n_node, dout, din); - const long total4 = n_node * dout / 4; - const bool residual = resnets[l] && dout == din; - auto launch = [&](auto act_tag) { - layer_epilogue_kernel - <<>>( - total4, dout, h, - bs[l].numel() ? bs[l].data_ptr() : nullptr, - idts[l].numel() ? idts[l].data_ptr() : nullptr, cur, - residual ? 1 : 0, h, adot); - }; - if (act == 0) { - launch(std::integral_constant{}); - } else { - launch(std::integral_constant{}); - } - FITTING_CHECK_LAUNCH("graph_fitting layer"); - cur = h; - din = dout; - } - head_kernel<<>>( - n_node, din, cur, w_head.data_ptr(), - b_head.numel() ? b_head.data_ptr() : nullptr, - bias_atom_e.data_ptr(), atype.data_ptr(), - e.data_ptr()); - FITTING_CHECK_LAUNCH("graph_fitting head"); + // Two-slot ping-pong for the activations: layer l writes slot ``l & 1`` + // while reading the previous layer's slot, so an activation is overwritten + // only after the next GEMM has consumed it (kernels run in stream order). A + // single-layer network requires only its output slot. + const int slots = plan.n_layer > 1 ? 2 : 1; + auto act_buf = torch::empty({slots, n_node, plan.width_max}, f32); + float* activation[2] = { + act_buf[0].data_ptr(), + slots > 1 ? act_buf[1].data_ptr() : act_buf[0].data_ptr()}; + fitting_forward_range( + at::cuda::getCurrentCUDAStream(), plan, x.data_ptr(), x.size(1), + atype.data_ptr(), ws, bs, resnets, w_head, b_head, bias_atom_e, act, + n_node, saved.data_ptr(), activation, e.data_ptr()); return {e, saved}; } -// Backward: d_x from the upstream d_e (fp64 (N, 1)). The saved derivative +// Backward: d_x from the upstream d_e (fp64 (N, 1)). The saved pre-activation // extent and fitting widths determine the output shape, so the descriptor is -// not retained solely for shape metadata. Two ping-pong dh buffers walk the -// layers from the head down. +// not retained solely for shape metadata. void graph_fitting_backward_core(torch::Tensor d_e, torch::Tensor saved, std::vector ws, + std::vector bs, std::vector resnets, torch::Tensor w_head, + int64_t act, torch::Tensor d_x) { - long total_width = 0; - for (const auto& weight : ws) { - total_width += weight.size(1); - } - TORCH_CHECK(total_width > 0 && saved.numel() % total_width == 0, - "graph_fitting_backward: saved derivative buffer does not match " - "the fitting widths"); - const long n_node = saved.numel() / total_width; + const FittingLayerPlan plan = fitting_layer_plan(ws); + TORCH_CHECK(plan.saved_width() > 0 && saved.numel() % plan.saved_width() == 0, + "graph_fitting_backward: saved buffer does not match the " + "fitting widths"); + const long n_node = saved.numel() / plan.saved_width(); const long input_width = ws[0].size(0); TORCH_CHECK(d_x.dim() == 2 && d_x.size(0) == n_node && d_x.size(1) == input_width && @@ -300,81 +442,129 @@ void graph_fitting_backward_core(torch::Tensor d_e, if (n_node == 0) { return; } - auto stream = at::cuda::getCurrentCUDAStream(); - const int n_layer = (int)ws.size(); - std::vector offset(n_layer + 1, 0); - for (int l = 0; l < n_layer; ++l) { - offset[l + 1] = offset[l] + ws[l].size(1); - } + const c10::cuda::CUDAGuard device_guard(saved.device()); auto f32 = saved.options().dtype(torch::kFloat32); auto d_e_c = d_e.contiguous(); - - long width_max = 0; - for (int l = 0; l < n_layer; ++l) { - width_max = std::max(width_max, (long)ws[l].size(1)); - } - auto dh = torch::empty({n_node, width_max}, f32); - auto dh_next = torch::empty({n_node, width_max}, f32); - - { - const int dout = (int)ws[n_layer - 1].size(1); - seed_kernel<<>>( - n_node * dout / 4, dout, d_e_c.data_ptr(), - w_head.data_ptr(), dh.data_ptr()); - FITTING_CHECK_LAUNCH("graph_fitting seed"); - } - for (int l = n_layer - 1; l >= 0; --l) { - const int dout = (int)ws[l].size(1); - const int din = (int)ws[l].size(0); - const float* adot = saved.data_ptr() + offset[l] * n_node; - float* out = l > 0 ? dh_next.data_ptr() : d_x.data_ptr(); - const bool residual = resnets[l] && dout == din; - float beta = 0.f; - if (residual) { - // Identity bypass: dh_{l-1} starts from dh_l. - cudaMemcpyAsync(out, dh.data_ptr(), sizeof(float) * n_node * din, - cudaMemcpyDeviceToDevice, stream); - beta = 1.f; - } - backward_epilogue_kernel<<>>( - n_node * dout / 4, dh.data_ptr(), adot, dh.data_ptr()); - FITTING_CHECK_LAUNCH("graph_fitting backward layer"); - gemm_nt(stream, dh.data_ptr(), ws[l].data_ptr(), out, - (int)n_node, din, dout, beta); - if (l > 0) { - std::swap(dh, dh_next); - } - } + auto dh = torch::empty({n_node, plan.width_max}, f32); + auto dh_next = plan.n_layer > 1 ? torch::empty({n_node, plan.width_max}, f32) + : torch::empty({0}, f32); + fitting_backward_range(at::cuda::getCurrentCUDAStream(), plan, + d_e_c.data_ptr(), saved.data_ptr(), ws, + bs, resnets, w_head, act, n_node, dh.data_ptr(), + plan.n_layer > 1 ? dh_next.data_ptr() : nullptr, + d_x.data_ptr()); } torch::Tensor graph_fitting_backward(torch::Tensor d_e, torch::Tensor saved, std::vector ws, + std::vector bs, std::vector resnets, - torch::Tensor w_head) { - long total_width = 0; - for (const auto& weight : ws) { - total_width += weight.size(1); - } - TORCH_CHECK(total_width > 0 && saved.numel() % total_width == 0, - "graph_fitting_backward: saved derivative buffer does not match " - "the fitting widths"); - const long n_node = saved.numel() / total_width; + torch::Tensor w_head, + int64_t act) { + const FittingLayerPlan plan = fitting_layer_plan(ws); + TORCH_CHECK(plan.saved_width() > 0 && saved.numel() % plan.saved_width() == 0, + "graph_fitting_backward: saved buffer does not match the " + "fitting widths"); + const long n_node = saved.numel() / plan.saved_width(); auto d_x = torch::empty({n_node, ws[0].size(0)}, saved.options()); - graph_fitting_backward_core(d_e, saved, std::move(ws), std::move(resnets), - w_head, d_x); + graph_fitting_backward_core(d_e, saved, std::move(ws), std::move(bs), + std::move(resnets), w_head, act, d_x); return d_x; } +// Energy and input gradient in one pass over node tiles. +// +// Inference knows the head cotangent before the forward runs -- it is the +// ownership mask -- so the two directions need not be separated by the whole +// node axis. Walking tiles instead retires each tile's pre-activations as soon +// as its backward consumes them, which replaces the largest node-scale +// allocation of a step with one of tile size. The network is evaluated per +// node, so no tile needs anything from another and nothing is recomputed. +// +// ``tile`` of zero, or any value covering the system, degenerates to a single +// run and reproduces the untiled sequence exactly. +// +// The cotangent replaces the descriptor in place. A tile reads its descriptor +// rows in the first layer and writes their cotangent in the last, and no later +// tile revisits them, so the two never need separate node-scale allocations; +// inference has no further use for the descriptor once the forward has read +// it. Only the energy is returned. +torch::Tensor graph_fitting_energy_gradient(torch::Tensor x, + torch::Tensor atype, + std::vector ws, + std::vector bs, + std::vector resnets, + torch::Tensor w_head, + torch::Tensor b_head, + torch::Tensor bias_atom_e, + int64_t act, + torch::Tensor seed, + int64_t tile) { + TORCH_CHECK( + x.is_cuda() && x.is_contiguous() && x.scalar_type() == torch::kFloat32, + "graph_fitting_energy_gradient: x must be contiguous CUDA fp32"); + const c10::cuda::CUDAGuard device_guard(x.device()); + const long n_node = x.size(0); + const long input_width = x.size(1); + const FittingLayerPlan plan = fitting_layer_plan(ws); + TORCH_CHECK(plan.n_layer > 0 && ws[0].size(0) == input_width, + "graph_fitting_energy_gradient: first weight does not match the " + "descriptor width"); + auto f32 = x.options().dtype(torch::kFloat32); + auto e = torch::empty({n_node, 1}, x.options().dtype(torch::kFloat64)); + if (n_node == 0) { + return e; + } + auto seed_c = seed.contiguous(); + TORCH_CHECK( + seed_c.numel() == n_node && seed_c.scalar_type() == torch::kFloat64, + "graph_fitting_energy_gradient: seed must be fp64 with one " + "entry per node"); + + const long run = tile > 0 ? std::min(tile, n_node) : n_node; + const int slots = plan.n_layer > 1 ? 2 : 1; + auto saved = torch::empty({run * plan.saved_width()}, f32); + // The backward reads only the saved pre-activations, never the activations, + // so the forward ping-pong and the two cotangent buffers never hold live + // data at the same time and share one allocation. + auto scratch = torch::empty({slots, run, plan.width_max}, f32); + float* slot[2] = { + scratch[0].data_ptr(), + slots > 1 ? scratch[1].data_ptr() : scratch[0].data_ptr()}; + + auto stream = at::cuda::getCurrentCUDAStream(); + for (long begin = 0; begin < n_node; begin += run) { + const long count = std::min(run, n_node - begin); + fitting_forward_range( + stream, plan, x.data_ptr() + begin * input_width, input_width, + atype.data_ptr() + begin, ws, bs, resnets, w_head, b_head, + bias_atom_e, act, count, saved.data_ptr(), slot, + e.data_ptr() + begin); + fitting_backward_range(stream, plan, seed_c.data_ptr() + begin, + saved.data_ptr(), ws, bs, resnets, w_head, + act, count, slot[0], + plan.n_layer > 1 ? slot[1] : nullptr, + x.data_ptr() + begin * input_width); + } + return e; +} + TORCH_LIBRARY_FRAGMENT(deepmd, m) { m.def( "graph_fitting(Tensor x, Tensor atype, Tensor[] ws, Tensor[] bs, " - "Tensor[] idts, int[] resnets, Tensor w_head, Tensor b_head, " - "Tensor bias_atom_e, int act) -> (Tensor e, Tensor saved)"); + "int[] resnets, Tensor w_head, Tensor b_head, Tensor bias_atom_e, " + "int act) -> (Tensor e, Tensor saved)"); m.impl("graph_fitting", torch::kCUDA, &graph_fitting); m.def( - "graph_fitting_backward(Tensor d_e, Tensor saved, " - "Tensor[] ws, int[] resnets, Tensor w_head) -> Tensor"); + "graph_fitting_backward(Tensor d_e, Tensor saved, Tensor[] ws, " + "Tensor[] bs, int[] resnets, Tensor w_head, int act) -> Tensor"); m.impl("graph_fitting_backward", torch::kCUDA, &graph_fitting_backward); + m.def( + "graph_fitting_energy_gradient(Tensor(a!) x, Tensor atype, " + "Tensor[] ws, Tensor[] bs, int[] resnets, Tensor w_head, " + "Tensor b_head, Tensor bias_atom_e, int act, Tensor seed, int tile) " + "-> Tensor"); + m.impl("graph_fitting_energy_gradient", torch::kCUDA, + &graph_fitting_energy_gradient); } diff --git a/source/op/pt/graph_ops.h b/source/op/pt/graph_ops.h index 418364ef0b..ee6de887dd 100644 --- a/source/op/pt/graph_ops.h +++ b/source/op/pt/graph_ops.h @@ -96,25 +96,74 @@ torch::Tensor dpa1_graph_descriptor_backward( double nnei); // Energy fitting MLP on the flat node axis. Returns (atom_energy (N, 1) fp64, -// saved activations/derivatives) for graph_fitting_backward. +// saved layer pre-activations) for graph_fitting_backward. std::tuple graph_fitting( torch::Tensor x, torch::Tensor atype, std::vector ws, std::vector bs, - std::vector idts, std::vector resnets, torch::Tensor w_head, torch::Tensor b_head, torch::Tensor bias_atom_e, int64_t act); -// dE/d(x) from dE/d(atom_energy); consumes the saved activations. +// dE/d(x) from dE/d(atom_energy); consumes the saved pre-activations. torch::Tensor graph_fitting_backward(torch::Tensor d_e, torch::Tensor saved, std::vector ws, + std::vector bs, std::vector resnets, - torch::Tensor w_head); + torch::Tensor w_head, + int64_t act); + +// Layer geometry of one fitting network, shared by the operators that +// evaluate it over a run of nodes. +struct FittingLayerPlan { + std::vector offset; //!< Prefix sum of the hidden widths. + long width_max; //!< Widest hidden layer. + int n_layer; + + long saved_width() const { return offset[n_layer]; } +}; + +FittingLayerPlan fitting_layer_plan(const std::vector& ws); + +// Evaluate the fitting network over one contiguous run of nodes. Every +// full-width pointer is already indexed from the run's first node, so the same +// code serves the whole node axis and a single tile of it. ``saved`` and +// ``activation`` are sized for the run rather than for the system. +void fitting_forward_range(cudaStream_t stream, + const FittingLayerPlan& plan, + const float* x, + long input_width, + const long* atype, + const std::vector& ws, + const std::vector& bs, + const std::vector& resnets, + const torch::Tensor& w_head, + const torch::Tensor& b_head, + const torch::Tensor& bias_atom_e, + int64_t act, + long run_nodes, + float* saved, + float* const activation[2], + double* e); + +// Propagate the head cotangent of one run of nodes back to the input. +void fitting_backward_range(cudaStream_t stream, + const FittingLayerPlan& plan, + const double* d_e, + const float* saved, + const std::vector& ws, + const std::vector& bs, + const std::vector& resnets, + const torch::Tensor& w_head, + int64_t act, + long run_nodes, + float* dh, + float* dh_next, + float* d_x); // Scatter dE/d(edge_vec) into per-node force, per-frame virial and (optional) // per-node virial. Returns (force (N, 3), atom_virial (N, 3, 3) or empty, diff --git a/source/tests/common/dpmodel/test_descriptor_dpa4c.py b/source/tests/common/dpmodel/test_descriptor_dpa4c.py new file mode 100644 index 0000000000..8ba2220b98 --- /dev/null +++ b/source/tests/common/dpmodel/test_descriptor_dpa4c.py @@ -0,0 +1,756 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import dataclasses +from typing import ( + Any, +) +from unittest import ( + mock, +) + +import numpy as np +import pytest + +from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, +) +from deepmd.dpmodel.descriptor.dpa4c_nn import ( + build_angular_basis, + build_bispectrum_layout, + enumerate_degree_triples, + packed_l2_to_stf, +) +from deepmd.dpmodel.utils import ( + neighbor_graph, +) +from deepmd.dpmodel.utils.lebedev import ( + load_lebedev_rule, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + graph_from_dense_quartet, +) +from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, +) +from deepmd.dpmodel.utils.update_sel import ( + UpdateSel, +) + +COORD = np.array( + [ + [ + [0.0, 0.0, 0.0], + [1.1, 0.2, -0.1], + [-0.4, 0.9, 0.3], + [0.2, -0.5, 1.2], + [-0.7, -0.3, -0.8], + ] + ], + dtype=np.float64, +) +ATYPE = np.array([[0, 1, 0, 1, 0]], dtype=np.int64) + + +def dense_inputs( + descriptor: DescrptDPA4C, + coord: np.ndarray = COORD, + atype: np.ndarray = ATYPE, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Build a complete bounded neighbor list for a small test system.""" + return extend_input_and_build_neighbor_list( + coord, + atype, + descriptor.get_rcut(), + [8], + mixed_types=True, + box=None, + ) + + +def evaluate( + descriptor: DescrptDPA4C, + coord: np.ndarray = COORD, + atype: np.ndarray = ATYPE, +) -> np.ndarray: + """Evaluate a descriptor through the dense compatibility interface.""" + coord_ext, atype_ext, mapping, nlist = dense_inputs( + descriptor, + coord, + atype, + ) + return descriptor( + coord_ext, + atype_ext, + nlist, + mapping=mapping, + )[0] + + +def build_graph( + descriptor: DescrptDPA4C, + coord: np.ndarray = COORD, + atype: np.ndarray = ATYPE, +) -> tuple[Any, np.ndarray]: + """Build the flat neighbor graph consumed by the graph-native equations.""" + coord_ext, atype_ext, mapping, nlist = dense_inputs(descriptor, coord, atype) + return graph_from_dense_quartet(coord_ext, atype_ext, nlist, mapping) + + +def edge_features( + descriptor: DescrptDPA4C, + graph: Any, + atype_local: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return the masked edge amplitudes, harmonics, and cutoff envelope.""" + return descriptor.build_edge_features( + graph, + atype_local, + *descriptor.pair_film.call(descriptor.type_embedding.call()), + ) + + +def moment_blocks( + descriptor: DescrptDPA4C, + moments: np.ndarray, +) -> list[np.ndarray]: + """Split flat moments and apply the fixed degree-one/two alignment.""" + readout = descriptor.readout + blocks = [] + for degree, width in enumerate(descriptor.degree_channels): + start, end = readout.degree_offsets[degree : degree + 2] + blocks.append( + moments[:, start:end].reshape( + moments.shape[0], + 2 * degree + 1, + width, + ) + ) + for degree, projection in enumerate(readout.channel_alignment, start=1): + blocks[degree] = projection.call(blocks[degree]) + return blocks + + +def projected_blocks( + descriptor: DescrptDPA4C, + moments: np.ndarray, +) -> list[np.ndarray]: + """Build effective low-rank blocks from aligned moments.""" + blocks = moment_blocks(descriptor, moments) + return [ + block if projection is None else projection.call(block) + for projection, block in zip( + descriptor.readout.probe_projections, + blocks[1:], + strict=True, + ) + ] + + +class TestDPA4C: + def setup_method(self) -> None: + self.descriptor = DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=8, + lmax=2, + n_radial=4, + precision="float64", + seed=17, + ) + + def test_single_reduction_dense_graph_parity( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + calls = 0 + segment_sum = neighbor_graph.segment_sum + + def count_segment_sum(*args: Any, **kwargs: Any) -> Any: + nonlocal calls + calls += 1 + return segment_sum(*args, **kwargs) + + monkeypatch.setattr(neighbor_graph, "segment_sum", count_segment_sum) + coord_ext, atype_ext, mapping, nlist = dense_inputs(self.descriptor) + dense = self.descriptor( + coord_ext, + atype_ext, + nlist, + mapping=mapping, + )[0] + assert calls == 1 + + graph, atype_local = graph_from_dense_quartet( + coord_ext, + atype_ext, + nlist, + mapping, + ) + calls = 0 + graph_output, rotation = self.descriptor.call_graph(graph, atype_local) + assert calls == 1 + assert rotation is None + np.testing.assert_allclose( + graph_output.reshape(dense.shape), + dense, + atol=1e-12, + rtol=1e-12, + ) + + def test_moments_match_explicit_reference(self) -> None: + descriptor = DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=32, + lmax=2, + n_radial=4, + radial_modes=2, + precision="float64", + seed=17, + ) + graph, atype_local = build_graph(descriptor) + amplitude, basis, envelope = edge_features(descriptor, graph, atype_local) + dst = graph.edge_index[1] + n_total = atype_local.shape[0] + moments, divisors = descriptor.aggregate_moments( + amplitude, + basis, + envelope, + dst, + n_total, + ) + scalar = moments[:, : descriptor.channels] + angular = moments[:, descriptor.channels :] + scalar_mass = np.zeros(n_total, dtype=amplitude.dtype) + angular_mass = np.zeros(n_total, dtype=amplitude.dtype) + np.add.at(scalar_mass, dst, envelope**2) + np.add.at(angular_mass, dst, envelope**4) + floor = descriptor._DEGREE_NORM_FLOOR + expected_scalar_normalizer = 1.0 / np.sqrt(scalar_mass + floor) + expected_angular_normalizer = 1.0 / np.sqrt(angular_mass + floor) + + # The reduction also returns the two divisors, which are exactly the + # reciprocals of the normalizers applied to the moments. + np.testing.assert_allclose( + divisors, + np.stack( + [1.0 / expected_scalar_normalizer, 1.0 / expected_angular_normalizer], + axis=-1, + ), + atol=2e-13, + rtol=2e-13, + ) + reduced_scalar = np.zeros( + (n_total, descriptor.channels), + dtype=amplitude.dtype, + ) + np.add.at(reduced_scalar, dst, amplitude[:, : descriptor.channels]) + np.testing.assert_allclose( + scalar, + reduced_scalar * expected_scalar_normalizer[:, None], + atol=2e-13, + rtol=2e-13, + ) + + # Every non-scalar degree carries one additional envelope factor and + # the matched normalizer derived from the fourth envelope power. + expected_blocks = [] + for degree in (1, 2): + channels = list(range(descriptor.degree_channels[degree])) + edge_value = ( + amplitude[:, channels][:, None, :] + * basis[:, degree**2 : (degree + 1) ** 2, None] + * envelope[:, None, None] + ) + reduced = np.zeros( + (n_total, 2 * degree + 1, len(channels)), + dtype=edge_value.dtype, + ) + np.add.at(reduced, dst, edge_value) + reduced *= expected_angular_normalizer[:, None, None] + expected_blocks.append(reduced.reshape(n_total, -1)) + np.testing.assert_allclose( + angular, + np.concatenate(expected_blocks, axis=1), + atol=2e-13, + rtol=2e-13, + ) + + def test_lmax4_o3_translation_and_permutation_invariance(self) -> None: + descriptor = DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=16, + lmax=4, + n_radial=4, + precision="float64", + seed=19, + ) + rotation = np.array( + [ + [-2.0 / 3.0, 2.0 / 15.0, 11.0 / 15.0], + [2.0 / 3.0, -1.0 / 3.0, 2.0 / 3.0], + [1.0 / 3.0, 14.0 / 15.0, 2.0 / 15.0], + ], + dtype=np.float64, + ) + + reference = evaluate(descriptor) + rotated = evaluate(descriptor, COORD @ rotation.T) + reflected_coord = COORD.copy() + reflected_coord[..., 0] *= -1.0 + reflected = evaluate(descriptor, reflected_coord) + translated = evaluate( + descriptor, + COORD + np.array([1.7, -0.8, 2.1]), + ) + permutation = np.array([2, 4, 0, 3, 1]) + permuted = evaluate( + descriptor, + COORD[:, permutation], + ATYPE[:, permutation], + ) + + np.testing.assert_allclose(rotated, reference, atol=2e-12, rtol=2e-12) + np.testing.assert_allclose(reflected, reference, atol=2e-12, rtol=2e-12) + np.testing.assert_allclose(translated, reference, atol=2e-12, rtol=2e-12) + np.testing.assert_allclose( + permuted, + reference[:, permutation], + atol=2e-12, + rtol=2e-12, + ) + + def test_serialization_roundtrip_preserves_fixed_structure(self) -> None: + self.descriptor.compute_input_stats([{"coord": COORD, "atype": ATYPE}]) + reference = evaluate(self.descriptor) + + # The two neighborhood masses are standardized rather than merely + # rescaled: their information lies on an offset far larger than any + # other invariant carries, so only these coordinates take a mean. + geometry_end = self.descriptor.get_dim_out() - self.descriptor.channels + mass = slice(geometry_end - 2, geometry_end) + assert np.all(self.descriptor.mean[mass] > 0.0) + np.testing.assert_array_equal(np.delete(self.descriptor.mean, mass), 0.0) + + data = self.descriptor.serialize() + assert data["radial_embedding"]["mlp_layers"] == [4, 24, 8] + assert data["channels"] == 8 + assert data["lmax"] == 2 + assert "bispectrum_ranks" not in data + assert data["readout"]["channels"] == 8 + assert data["readout"]["lmax"] == 2 + assert "bispectrum_ranks" not in data["readout"] + + restored = DescrptDPA4C.deserialize(data) + result = evaluate(restored) + np.testing.assert_array_equal(result, reference) + np.testing.assert_array_equal(restored.mean, self.descriptor.mean) + np.testing.assert_array_equal(restored.stddev, self.descriptor.stddev) + + default_data = DescrptDPA4C(rcut=3.0, ntypes=1, use_amp=True).serialize() + assert default_data["channels"] == 32 + assert default_data["lmax"] == 2 + assert default_data["radial_modes"] == 0 + # Mixed precision is an execution policy supplied at load time, so a + # checkpoint must not carry it. + assert "use_amp" not in default_data + + @pytest.mark.parametrize( + "divergence", + [ + {"rcut": 4.0}, + {"ntypes": 3}, + {"channels": 16}, + {"lmax": 3}, + {"basis_type": "gaussian"}, + {"n_radial": 8}, + {"radial_modes": 2}, + {"use_amp": True}, + {"trainable": False}, + {"type_map": ["H", "O"]}, + {"precision": "float32"}, + ], + ) + def test_sharing_rejects_incompatible_structures( + self, + divergence: dict[str, Any], + ) -> None: + """Every field that shapes a shared module must block sharing. + + Sharing binds the type table, radial basis, radial network, mode head, + pair cache, and readout of the base descriptor into the replica. A + divergence that the signature fails to catch is silent rather than + loud: gathers may run out of bounds, a replica may inherit frozen + weights while still reporting itself trainable, or two branches may + read the same type table under different element orders. + """ + config = { + "rcut": 3.0, + "ntypes": 2, + "channels": 8, + "lmax": 2, + "basis_type": "bessel", + "n_radial": 4, + "radial_modes": 1, + "trainable": True, + "type_map": ["O", "H"], + "precision": "float64", + } + base = DescrptDPA4C(**config, seed=0) + DescrptDPA4C(**config, seed=1).share_params(base, 0) + with pytest.raises(ValueError, match="identical structural parameters"): + DescrptDPA4C(**{**config, **divergence}, seed=1).share_params(base, 0) + + def test_shared_replica_matches_the_base_descriptor(self) -> None: + config = { + "rcut": 3.0, + "ntypes": 2, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "radial_modes": 1, + "precision": "float64", + } + base = DescrptDPA4C(**config, seed=0) + replica = DescrptDPA4C(**config, seed=1) + assert not np.allclose(evaluate(replica), evaluate(base)) + replica.share_params(base, 0) + np.testing.assert_array_equal(evaluate(replica), evaluate(base)) + + def test_aligned_grams_match_explicit_contractions(self) -> None: + descriptor = DescrptDPA4C( + rcut=3.0, + ntypes=1, + channels=16, + lmax=4, + precision="float64", + seed=29, + ) + rng = np.random.default_rng(31) + moments = rng.normal(size=(5, descriptor.readout.degree_offsets[-1])) + output = descriptor.readout.call(moments) + blocks = moment_blocks(descriptor, moments) + + cursor = descriptor.channels + for block, width in zip( + blocks[1:], + descriptor.degree_channels[1:], + strict=True, + ): + gram = np.transpose(block, (0, 2, 1)) @ block + row, column = np.triu_indices(width) + expected = gram[:, row, column] + expected *= np.where(row == column, 1.0, np.sqrt(2.0))[None, :] + actual = output[:, cursor : cursor + row.size] + np.testing.assert_allclose(actual, expected, atol=1e-13, rtol=1e-13) + # The off-diagonal scale exists to make the half-vectorization + # norm preserving, so assert that property rather than the + # constant the implementation happens to use. + np.testing.assert_allclose( + np.sum(actual**2, axis=1), + np.sum(gram**2, axis=(1, 2)), + atol=1e-12, + rtol=1e-12, + ) + cursor += row.size + assert cursor == descriptor.channels + descriptor.readout.gram_index.size + + def test_bispectrum_and_quartic_match_explicit_references( + self, + ) -> None: + descriptor = DescrptDPA4C( + rcut=3.0, + ntypes=1, + channels=32, + lmax=4, + precision="float64", + seed=37, + ) + rng = np.random.default_rng(37) + moments = rng.normal(size=(3, descriptor.readout.degree_offsets[-1])) + output = descriptor.readout.call(moments) + projected = projected_blocks(descriptor, moments) + + expected_parts = [] + for triple_index, degrees in enumerate(descriptor.readout.degree_triples): + coupling_start, coupling_end = descriptor.readout.coupling_offsets[ + triple_index : triple_index + 2 + ] + coupling = descriptor.readout.bispectrum_coupling[ + coupling_start:coupling_end + ].reshape(*(2 * degree + 1 for degree in degrees)) + full = np.einsum( + "ijk,nia,njb,nkc->nabc", + coupling, + projected[degrees[0] - 1], + projected[degrees[1] - 1], + projected[degrees[2] - 1], + optimize=True, + ).reshape(3, -1) + probe_start, probe_end = descriptor.readout.probe_offsets[ + triple_index : triple_index + 2 + ] + reduced = ( + full[ + :, + descriptor.readout.probe_index[probe_start:probe_end], + ] + * descriptor.readout.probe_scale[None, probe_start:probe_end] + ) + # Equal-degree axes emit one representative per orbit. The + # multiplicity scale exists so that dropping the rest preserves + # the norm of the full symmetric tensor. + np.testing.assert_allclose( + np.sum(reduced**2, axis=1), + np.sum(full**2, axis=1), + atol=1e-12, + rtol=1e-12, + ) + expected_parts.append(reduced) + expected_bispectrum = np.concatenate(expected_parts, axis=1) + bispectrum_start = descriptor.channels + descriptor.readout.gram_index.size + bispectrum_end = bispectrum_start + expected_bispectrum.shape[1] + np.testing.assert_allclose( + output[:, bispectrum_start:bispectrum_end], + expected_bispectrum, + atol=2e-13, + rtol=2e-13, + ) + + vectors = np.transpose(projected[0], (0, 2, 1)) + tensors = packed_l2_to_stf(np.transpose(projected[1], (0, 2, 1))) + tensor_vector = np.matmul( + tensors[:, :, None, :, :], + vectors[:, None, :, :, None], + )[..., 0] + expected_quartic = np.sum( + tensor_vector * tensor_vector, + axis=-1, + ).reshape(3, -1) + np.testing.assert_allclose( + output[:, bispectrum_end:], + expected_quartic, + atol=2e-13, + rtol=2e-13, + ) + + def test_jax_lmax4_matches_numpy(self) -> None: + pytest.importorskip("jax") + from deepmd.jax.env import ( + jax, + jnp, + ) + + descriptor = DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=16, + lmax=4, + n_radial=4, + radial_modes=2, + precision="float64", + seed=43, + ) + graph, atype_local = build_graph(descriptor) + reference, _ = descriptor.call_graph(graph, atype_local) + graph_jax = dataclasses.replace( + graph, + n_node=jnp.asarray(graph.n_node), + edge_index=jnp.asarray(graph.edge_index), + edge_vec=jnp.asarray(graph.edge_vec), + edge_mask=jnp.asarray(graph.edge_mask), + ) + atype_jax = jnp.asarray(atype_local) + + def evaluate_jax(edge_vec: object) -> object: + current_graph = dataclasses.replace(graph_jax, edge_vec=edge_vec) + return descriptor.call_graph(current_graph, atype_jax)[0] + + np.testing.assert_allclose( + np.asarray(jax.jit(evaluate_jax)(graph_jax.edge_vec)), + reference, + atol=3e-10, + rtol=3e-10, + ) + + def test_addition_theorem_and_parity_through_degree_four(self) -> None: + rng = np.random.default_rng(41) + left = rng.normal(size=(32, 3)) + right = rng.normal(size=(32, 3)) + left /= np.linalg.norm(left, axis=-1, keepdims=True) + right /= np.linalg.norm(right, axis=-1, keepdims=True) + left_basis = build_angular_basis(left, 4) + right_basis = build_angular_basis(right, 4) + reflected_basis = build_angular_basis(-left, 4) + cosine = np.sum(left * right, axis=-1) + legendre = ( + np.ones_like(cosine), + cosine, + 0.5 * (3.0 * cosine**2 - 1.0), + 0.5 * (5.0 * cosine**3 - 3.0 * cosine), + 0.125 * (35.0 * cosine**4 - 30.0 * cosine**2 + 3.0), + ) + + for degree in range(5): + start, end = degree**2, (degree + 1) ** 2 + np.testing.assert_allclose( + np.sum( + left_basis[:, start:end] * right_basis[:, start:end], + axis=1, + ), + legendre[degree], + atol=2e-15, + rtol=2e-15, + ) + np.testing.assert_allclose( + reflected_basis[:, start:end], + (-1) ** degree * left_basis[:, start:end], + atol=2e-15, + rtol=2e-15, + ) + + def test_allowed_degree_triples_and_gaunt_couplings(self) -> None: + expected_triples = ( + (1, 1, 2), + (1, 2, 3), + (1, 3, 4), + (2, 2, 2), + (2, 2, 4), + (2, 3, 3), + (2, 4, 4), + (3, 3, 4), + (4, 4, 4), + ) + assert enumerate_degree_triples(4) == expected_triples + + layout = build_bispectrum_layout(4, [1, 1, 1, 1]) + points, weights = load_lebedev_rule(17) + basis = build_angular_basis(points, 4) + for triple_index, degrees in enumerate(expected_triples): + start, end = layout.coupling_offsets[triple_index : triple_index + 2] + coupling = layout.coupling[start:end].reshape( + *(2 * degree + 1 for degree in degrees) + ) + degree_1, degree_2, degree_3 = degrees + reference = np.einsum( + "n,ni,nj,nk->ijk", + weights, + basis[:, degree_1**2 : (degree_1 + 1) ** 2], + basis[:, degree_2**2 : (degree_2 + 1) ** 2], + basis[:, degree_3**2 : (degree_3 + 1) ** 2], + optimize=True, + ) + reference /= np.linalg.norm(reference) + first = np.flatnonzero(np.abs(reference) > 1.0e-14)[0] + if reference.flat[first] < 0.0: + reference = -reference + np.testing.assert_allclose(coupling, reference, atol=2e-14, rtol=2e-14) + np.testing.assert_allclose(np.linalg.norm(coupling), 1.0, atol=1e-15) + + +@pytest.mark.parametrize( + ("config", "error"), + [ + ({"rcut": 0.0}, ValueError), + ({"ntypes": 0}, ValueError), + ({"n_radial": 0}, ValueError), + ({"channels": 15}, ValueError), + ({"channels": 256}, ValueError), + ({"channels": 32.0}, TypeError), + ({"lmax": 1}, ValueError), + ({"lmax": 5}, ValueError), + ({"lmax": 2.0}, TypeError), + # `bool` is an `int` subclass, so the guards reject it explicitly. + ({"channels": True}, TypeError), + ({"radial_modes": True}, ValueError), + ({"radial_modes": -1}, ValueError), + ({"spin": {}}, NotImplementedError), + ], +) +def test_configuration_boundaries( + config: dict[str, Any], + error: type[Exception], +) -> None: + with pytest.raises(error): + DescrptDPA4C(**{"rcut": 3.0, "ntypes": 1, **config}) + + +@pytest.mark.parametrize("neighbors", [8, 4096]) +def test_neighbor_statistics_report_distance_without_deriving_a_capacity( + neighbors: int, +) -> None: + """Statistics must never reject an environment or introduce a ``sel``. + + The descriptor is graph-native, so an environment denser than the dense + adapter bound is still valid; only the minimum neighbor distance is read. + """ + config = {"type": "dpa4c", "rcut": 6.0, "channels": 32, "lmax": 2} + with mock.patch.object( + UpdateSel, + "get_nbor_stat", + return_value=(0.8, [neighbors]), + ): + updated, min_nbor_dist = DescrptDPA4C.update_sel(None, ["O", "H"], config) + + assert updated == config + assert "sel" not in updated + assert min_nbor_dist == 0.8 + + +def _bispectrum_dimension(lmax: int, ranks: list[int]) -> int: + """Return the independent bispectrum width from rank combinatorics.""" + dimension = 0 + for degree_1, degree_2, degree_3 in enumerate_degree_triples(lmax): + rank_1 = ranks[degree_1 - 1] + rank_2 = ranks[degree_2 - 1] + rank_3 = ranks[degree_3 - 1] + if degree_1 == degree_3: + dimension += rank_1 * (rank_1 + 1) * (rank_1 + 2) // 6 + elif degree_1 == degree_2: + dimension += rank_1 * (rank_1 + 1) * rank_3 // 2 + elif degree_2 == degree_3: + dimension += rank_1 * rank_2 * (rank_2 + 1) // 2 + else: + dimension += rank_1 * rank_2 * rank_3 + return dimension + + +@pytest.mark.parametrize( + ("channels", "base_degree_channels", "base_ranks"), + [ + (8, [8, 4, 4], [4, 2]), + (16, [16, 4, 4], [4, 2]), + (32, [32, 8, 4], [4, 2]), + (64, [64, 8, 4], [4, 2]), + (128, [128, 16, 8], [8, 2]), + ], +) +def test_automatic_profiles_and_output_dimensions( + channels: int, + base_degree_channels: list[int], + base_ranks: list[int], +) -> None: + for lmax in (2, 3, 4): + descriptor = DescrptDPA4C( + rcut=3.0, + ntypes=1, + channels=channels, + lmax=lmax, + radial_modes=3, + ) + degree_channels = base_degree_channels + [1] * (lmax - 2) + ranks = base_ranks + [1] * (lmax - 2) + assert descriptor.degree_channels == degree_channels + assert descriptor.bispectrum_ranks == ranks + + # The radial function class never enters the descriptor layout. The + # trailing pair is the two neighborhood masses. + expected_dim = ( + 2 * channels + + sum(width * (width + 1) // 2 for width in degree_channels[1:]) + + _bispectrum_dimension(lmax, ranks) + + ranks[0] * ranks[1] + + 2 + ) + assert descriptor.get_dim_out() == expected_dim diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 111436fd3b..7036c5f811 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -65,6 +65,7 @@ p_examples / "water" / "dpa4" / "input.json", p_examples / "water" / "dpa4" / "input-zbl.json", p_examples / "water" / "dpa4" / "lmp" / "input.json", + p_examples / "water" / "dpa4c" / "input.json", p_examples / "property" / "train" / "input_torch.json", p_examples / "water" / "se_e3_tebd" / "input_torch.json", p_examples / "hessian" / "single_task" / "input.json", diff --git a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py index e14dc81610..7280f66343 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py @@ -1126,8 +1126,37 @@ def run(level): def test_parity_plain(self) -> None: self._assert_parity(self._build(resnet_dt=False)) - def test_parity_resnet_dt_silu(self) -> None: - self._assert_parity(self._build(resnet_dt=True, act="silu")) + def test_parity_silu(self) -> None: + self._assert_parity(self._build(resnet_dt=False, act="silu")) + + def test_parity_single_layer_residual(self) -> None: + self._assert_parity( + self._build( + resnet_dt=False, + act="silu", + neuron=[64], + ) + ) + + def test_timestep_falls_back(self) -> None: + from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + ) + + self.assertFalse(fitting_eligible(self._build(resnet_dt=True))) + + def test_ineligible_network_is_refused_not_approximated(self) -> None: + """An unsupported network must raise where its arguments are built. + + The operator has no representation for a layer timestep and would + evaluate the network without it, so the conversion refuses instead. + """ + from deepmd.kernels.cuda.graph_fitting import ( + fitting_operator_arguments, + ) + + with self.assertRaises(ValueError): + fitting_operator_arguments(self._build(resnet_dt=True)) def test_fparam_falls_back(self) -> None: from deepmd.kernels.cuda.graph_fitting import ( @@ -1202,6 +1231,7 @@ def _build_fitting(self, dim_descrpt): ntypes=2, dim_descrpt=dim_descrpt, neuron=[48, 48], + resnet_dt=False, activation_function="silu", precision="float32", mixed_types=True, @@ -1504,7 +1534,7 @@ def _build_fitting(self, dim_descrpt, ntypes=2): ntypes=ntypes, dim_descrpt=dim_descrpt, neuron=[64, 64, 64], - resnet_dt=True, + resnet_dt=False, activation_function="silu", precision="float32", mixed_types=True, @@ -1705,6 +1735,7 @@ def test_compact_canonical_torch_export_contract(self) -> None: self.assertEqual(metadata["lower_input_kind"], "dpa1_canonical") self.assertEqual(metadata["graph_edge_dtype"], "float32") + self.assertEqual(metadata["canonical_index_dtype"], "uint32") self.assertNotIn("graph_index_dtype", metadata) self.assertEqual( output_keys, diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py new file mode 100644 index 0000000000..a3488e1cea --- /dev/null +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -0,0 +1,428 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import dataclasses +import os +from typing import ( + Any, +) +from unittest import ( + mock, +) + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.descriptor.dpa4c import DescrptDPA4C as DPDescrptDPA4C +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + graph_from_dense_quartet, +) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) +from deepmd.pt_expt.descriptor.dpa4c import ( + DescrptDPA4C, +) +from deepmd.pt_expt.utils import ( + env, +) + +# Structural variants exercised by the backend-agnostic contracts. They span +# both angular profiles and both radial function classes. +STRUCTURES = [ + {"channels": 16, "lmax": 2, "radial_modes": 0}, + {"channels": 32, "lmax": 4, "radial_modes": 0}, + {"channels": 32, "lmax": 2, "radial_modes": 3}, + {"channels": 16, "lmax": 3, "radial_modes": 3}, +] + + +class TestDPA4C: + def setup_method(self) -> None: + self.descriptor = self.build() + self.coord = torch.tensor( + [ + [ + [0.0, 0.0, 0.0], + [1.1, 0.2, -0.1], + [-0.4, 0.9, 0.3], + [0.2, -0.5, 1.2], + [-0.7, -0.3, -0.8], + ] + ], + dtype=torch.float64, + device=env.DEVICE, + ) + self.atype = torch.tensor( + [[0, 1, 0, 1, 0]], + dtype=torch.long, + device=env.DEVICE, + ) + + @staticmethod + def build(**structure: Any) -> DescrptDPA4C: + """Build a descriptor on the backend device. + + The default precision is double, which the numerical contracts need. + Mixed-precision tests must override it, because CUDA autocast ignores + double operands and would otherwise leave the region untouched. + """ + return DescrptDPA4C( + rcut=3.0, + ntypes=2, + n_radial=4, + seed=17, + **{ + "channels": 16, + "lmax": 2, + "precision": "float64", + **structure, + }, + ).to(env.DEVICE) + + def _inputs( + self, + coord: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return extend_input_and_build_neighbor_list( + coord, + self.atype, + self.descriptor.get_rcut(), + [8], + mixed_types=True, + box=None, + ) + + def _evaluate( + self, + descriptor: DescrptDPA4C, + coord: torch.Tensor, + ) -> torch.Tensor: + coord_ext, atype_ext, mapping, nlist = self._inputs(coord) + return descriptor(coord_ext, atype_ext, nlist, mapping=mapping)[0] + + def _dimer_probe( + self, + descriptor: DescrptDPA4C, + distance: torch.Tensor, + *, + active: bool = True, + ) -> torch.Tensor: + """Contract the descriptor of one undirected dimer into a scalar.""" + zero = torch.zeros_like(distance) + edge_vec = torch.stack( + [ + torch.stack([distance, zero, zero]), + torch.stack([-distance, zero, zero]), + ] + ) + graph = NeighborGraph( + n_node=torch.tensor([2], dtype=torch.long, device=env.DEVICE), + edge_index=torch.tensor( + [[1, 0], [0, 1]], + dtype=torch.long, + device=env.DEVICE, + ), + edge_vec=edge_vec, + edge_mask=torch.full( + (2,), + active, + dtype=torch.bool, + device=env.DEVICE, + ), + ) + atype = torch.zeros(2, dtype=torch.long, device=env.DEVICE) + output, _ = descriptor.call_graph(graph, atype) + cotangent = torch.linspace( + -0.7, + 1.3, + output.numel(), + dtype=output.dtype, + device=output.device, + ).reshape_as(output) + return (output * cotangent).sum() + + def _dimer_derivatives( + self, + descriptor: DescrptDPA4C, + distance: float, + ) -> list[torch.Tensor]: + """Return the value and first three radial derivatives of the probe.""" + radius = torch.tensor( + distance, + dtype=torch.float64, + device=env.DEVICE, + requires_grad=True, + ) + derivatives = [self._dimer_probe(descriptor, radius)] + for _ in range(3): + derivatives.append( + torch.autograd.grad( + derivatives[-1], + radius, + create_graph=True, + )[0] + ) + return derivatives + + # === Backend equivalence === + + @pytest.mark.parametrize("structure", STRUCTURES) + def test_numpy_pytorch_parity_and_parameter_gradients( + self, + structure: dict[str, int], + ) -> None: + descriptor = self.build(**structure) + coord_ext, atype_ext, mapping, nlist = self._inputs(self.coord) + result = descriptor(coord_ext, atype_ext, nlist, mapping=mapping)[0] + reference = DPDescrptDPA4C.deserialize(descriptor.serialize())( + coord_ext.cpu().numpy(), + atype_ext.cpu().numpy(), + nlist.cpu().numpy(), + mapping=mapping.cpu().numpy(), + )[0] + np.testing.assert_allclose( + result.detach().cpu().numpy(), + reference, + atol=3e-12, + rtol=3e-12, + ) + + # Every trainable array must remain reachable from the loss. + result.square().mean().backward() + for name, parameter in descriptor.named_parameters(): + assert parameter.grad is not None, name + assert torch.isfinite(parameter.grad).all(), name + + def test_sharing_keeps_the_branch_local_exclusion_mask(self) -> None: + """Sharing must not carry the pair-exclusion mask across replicas. + + ``exclude_types`` is absent from the structural signature because each + branch of a multitask model configures it separately. The mask is a + submodule, so a backend that rebinds the whole submodule table would + capture it and silently make the replica evaluate excluded pairs. + """ + base = self.build(exclude_types=[]) + replica = self.build(exclude_types=[[0, 1]]) + expected = replica.emask.type_mask.detach().clone() + replica.share_params(base, 0) + torch.testing.assert_close(replica.emask.type_mask, expected) + assert replica.exclude_types == [[0, 1]] + assert replica.readout is base.readout + + def test_serialization_preserves_parameters(self) -> None: + restored = DescrptDPA4C.deserialize(self.descriptor.serialize()).to(env.DEVICE) + original_parameters = dict(self.descriptor.named_parameters()) + restored_parameters = dict(restored.named_parameters()) + assert original_parameters.keys() == restored_parameters.keys() + for name in original_parameters: + torch.testing.assert_close( + restored_parameters[name], + original_parameters[name], + ) + + @pytest.mark.parametrize("structure", STRUCTURES) + def test_torch_export_matches_eager(self, structure: dict[str, int]) -> None: + descriptor = self.build(**structure).eval() + coord_ext, atype_ext, mapping, nlist = self._inputs(self.coord) + exported = torch.export.export( + descriptor, + (coord_ext, atype_ext, nlist), + kwargs={"mapping": mapping}, + strict=False, + ) + torch.testing.assert_close( + exported.module()(coord_ext, atype_ext, nlist, mapping=mapping)[0], + descriptor(coord_ext, atype_ext, nlist, mapping=mapping)[0], + ) + + # === Differentiability === + + def test_coordinate_gradient_matches_finite_difference(self) -> None: + coord = self.coord.detach().clone().requires_grad_(True) + output = self._evaluate(self.descriptor, coord) + cotangent = torch.linspace( + -0.8, + 1.1, + output.numel(), + dtype=output.dtype, + device=output.device, + ).reshape_as(output) + (gradient,) = torch.autograd.grad((output * cotangent).sum(), coord) + + epsilon = 1e-6 + finite_difference = torch.empty_like(coord) + flat = coord.detach().reshape(-1) + for index in range(flat.numel()): + shifted = [] + for sign in (1.0, -1.0): + probe = flat.clone() + probe[index] += sign * epsilon + shifted.append( + ( + self._evaluate(self.descriptor, probe.reshape_as(coord)) + * cotangent + ).sum() + ) + finite_difference.reshape(-1)[index] = (shifted[0] - shifted[1]) / ( + 2.0 * epsilon + ) + torch.testing.assert_close( + gradient, + finite_difference, + atol=3e-8, + rtol=3e-8, + ) + + @pytest.mark.parametrize("structure", STRUCTURES) + def test_force_loss_double_backward(self, structure: dict[str, int]) -> None: + descriptor = self.build(**structure) + coord = self.coord.detach().clone().requires_grad_(True) + output = self._evaluate(descriptor, coord) + (force_gradient,) = torch.autograd.grad( + output.square().sum(), + coord, + create_graph=True, + ) + named_parameters = list(descriptor.named_parameters()) + parameter_gradients = torch.autograd.grad( + force_gradient.square().mean(), + tuple(parameter for _, parameter in named_parameters), + allow_unused=True, + ) + for (name, _), gradient in zip( + named_parameters, + parameter_gradients, + strict=True, + ): + assert gradient is not None, name + assert torch.isfinite(gradient).all(), name + + # === Mixed precision === + + @pytest.mark.skipif( + env.DEVICE.type != "cuda", + reason="autocast only engages on CUDA", + ) + @pytest.mark.parametrize("training", [True, False]) + @pytest.mark.parametrize("amp_infer", [False, True]) + @pytest.mark.parametrize("use_amp", [False, True]) + def test_amp_spans_the_edge_stage_and_stops_at_its_boundary( + self, + use_amp: bool, + amp_infer: bool, + training: bool, + ) -> None: + """Autocast must cover the per-edge stage exactly, and only when asked. + + Training follows ``use_amp`` and evaluation follows ``DP_AMP_INFER``; + the two are independent because mixed precision at inference is a + throughput choice that must not require the model to have been trained + with it. Where autocast does engage it has to survive both the radial + trunk and the mode head, which restore the dtype of their input unless + the descriptor opts out, and it must not escape into the destination + reduction or the readout. + + The precision is single because CUDA autocast ignores double operands, + which would make every assertion below vacuous. + """ + with mock.patch.dict( + os.environ, + {"DP_AMP_INFER": "1" if amp_infer else "0"}, + clear=False, + ): + descriptor = self.build( + channels=32, + radial_modes=4, + precision="float32", + use_amp=use_amp, + ) + descriptor.train(training) + coord_ext, atype_ext, mapping, nlist = self._inputs(self.coord) + graph, atype_local = graph_from_dense_quartet( + coord_ext, + atype_ext, + nlist, + mapping, + ) + graph = dataclasses.replace( + graph, + edge_vec=graph.edge_vec.to(torch.float32), + ) + + observed: list[torch.dtype] = [] + handles = [ + layer.register_forward_hook( + lambda module, args, output: observed.append(output.dtype) + ) + for layer in ( + *descriptor.radial_embedding.layers, + descriptor.radial_mode_head, + ) + ] + try: + features = descriptor.build_edge_features( + graph, + atype_local, + *descriptor.pair_film.call(descriptor.type_embedding.call()), + ) + finally: + for handle in handles: + handle.remove() + + active = use_amp if training else amp_infer + expected = torch.bfloat16 if active else torch.float32 + assert observed and all(dtype is expected for dtype in observed) + for feature in features: + assert feature.dtype == torch.float32 + assert torch.isfinite(feature).all() + + # === Cutoff and regularization === + + def test_cutoff_is_c3_continuous(self) -> None: + descriptor = self.build(channels=16, lmax=4).eval() + inside = self._dimer_derivatives(descriptor, descriptor.rcut - 1.0e-5) + boundary = self._dimer_derivatives(descriptor, descriptor.rcut) + outside = self._dimer_derivatives(descriptor, descriptor.rcut + 1.0e-5) + for order in range(4): + torch.testing.assert_close( + boundary[order], + outside[order], + atol=1e-14, + rtol=0.0, + ) + # The value and its first three derivatives approach the cutoff from + # inside at the rate set by the p=5 envelope. + assert inside[1].abs() < 1e-12 + assert inside[2].abs() < 2e-8 + assert inside[3].abs() < 5e-4 + + def test_cutoff_edge_matches_removed_topology(self) -> None: + descriptor = self.build().eval() + radius = torch.tensor( + descriptor.rcut, + dtype=torch.float64, + device=env.DEVICE, + ) + torch.testing.assert_close( + self._dimer_probe(descriptor, radius, active=True), + self._dimer_probe(descriptor, radius, active=False), + atol=1e-14, + rtol=0.0, + ) + + def test_coincident_edge_has_finite_third_derivative(self) -> None: + descriptor = self.build(lmax=3).eval() + derivatives = self._dimer_derivatives(descriptor, 0.0) + for derivative in derivatives: + assert torch.isfinite(derivative) + # Direction regularization makes the probe even in the separation, so + # every odd radial derivative vanishes at coincidence. + for order in (1, 3): + torch.testing.assert_close( + derivatives[order], + torch.zeros_like(derivatives[order]), + atol=1e-12, + rtol=0.0, + ) diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py new file mode 100644 index 0000000000..ef5e5bf4b2 --- /dev/null +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py @@ -0,0 +1,1059 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Numerical contract of the compressed DPA4C CUDA mega kernel.""" + +import dataclasses + +import pytest +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + attach_edge_csr, + graph_from_dense_quartet, +) +from deepmd.kernels.cuda.dpa4c.graph_compress import ( + _cpu_descriptor, + _table_lookup, + build_compression_artifacts, + build_radial_table, + descriptor_profile, + dpa4c_graph_compress_energy_force, + ensure_registered, + op_available, +) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) +from deepmd.pt_expt.descriptor.dpa4c import ( + DescrptDPA4C, +) + +_GPU = pytest.mark.skipif( + not torch.cuda.is_available() or not op_available(), + reason="CUDA and the compiled DPA4C operator are required", +) + + +def _build_descriptor( + channels: int, + lmax: int = 2, + radial_modes: int = 0, +) -> DescrptDPA4C: + return ( + DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=channels, + lmax=lmax, + n_radial=8, + radial_modes=radial_modes, + precision="float32", + seed=17, + ) + .cuda() + .eval() + ) + + +def _build_graph( + descriptor: DescrptDPA4C, + canonical: bool, + node_count: int = 24, +): + generator = torch.Generator(device="cuda").manual_seed(23) + coordinate = torch.rand( + 1, + node_count, + 3, + dtype=torch.float32, + device="cuda", + generator=generator, + ) + coordinate = coordinate * 5.0 + atype = torch.arange(node_count, device="cuda").reshape(1, -1) % 2 + coord_ext, atype_ext, mapping, nlist = extend_input_and_build_neighbor_list( + coordinate, + atype, + descriptor.rcut, + [48], + mixed_types=True, + box=None, + ) + graph, flat_type = graph_from_dense_quartet( + coord_ext, + atype_ext, + nlist, + mapping, + ) + graph = attach_edge_csr( + graph, + flat_type.shape[0], + canonicalize=canonical, + ) + return graph, flat_type + + +def _arguments( + descriptor: DescrptDPA4C, + graph, + atype: torch.Tensor, +): + ensure_registered() + artifacts = build_compression_artifacts(descriptor) + return ( + graph.edge_index, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + atype, + artifacts["data"], + artifacts["pair_film"], + artifacts["pair_mixing"], + artifacts["type_embedding"], + artifacts["readout_matrices"], + artifacts["coupling_meta"], + artifacts["coupling_entry"], + artifacts["coupling_value"], + artifacts["output_mean"], + artifacts["output_inv_std"], + bool(graph.destination_sorted), + int(descriptor.lmax), + *(float(value) for value in artifacts["info"]), + ) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 16, 32, 64, 128]) +@pytest.mark.parametrize("canonical", [False, True]) +def test_forward_backward_parity(channels: int, canonical: bool) -> None: + descriptor = _build_descriptor(channels) + graph, atype = _build_graph(descriptor, canonical) + arguments = _arguments(descriptor, graph, atype) + ensure_registered() + + edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) + output, state = torch.ops.deepmd.dpa4c_graph_compress( + edge_vec, + *arguments, + ) + assert state.shape == ( + atype.shape[0], + descriptor_profile(channels, descriptor.lmax).state_width, + ) + cotangent = torch.linspace( + -0.7, + 1.3, + output.numel(), + dtype=output.dtype, + device=output.device, + ).reshape_as(output) + (gradient,) = torch.autograd.grad((output * cotangent).sum(), edge_vec) + + reference_edge = graph.edge_vec.detach().clone().requires_grad_(True) + reference = _cpu_descriptor(reference_edge, *arguments) + (reference_gradient,) = torch.autograd.grad( + (reference * cotangent).sum(), + reference_edge, + ) + torch.testing.assert_close(output, reference, atol=2e-6, rtol=2e-6) + torch.testing.assert_close( + gradient, + reference_gradient, + atol=8e-6 if channels >= 64 else 3e-6, + rtol=1e-4 if channels >= 64 else 3e-5, + ) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 64, 128]) +def test_backward_tail_node_groups(channels: int) -> None: + descriptor = _build_descriptor(channels) + graph, atype = _build_graph(descriptor, canonical=True, node_count=23) + arguments = _arguments(descriptor, graph, atype) + cotangent = torch.randn( + atype.shape[0], + descriptor.get_dim_out(), + dtype=torch.float32, + device="cuda", + ) + + edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) + output, _state = torch.ops.deepmd.dpa4c_graph_compress( + edge_vec, + *arguments, + ) + (gradient,) = torch.autograd.grad((output * cotangent).sum(), edge_vec) + + reference_edge = graph.edge_vec.detach().clone().requires_grad_(True) + reference = _cpu_descriptor(reference_edge, *arguments) + (reference_gradient,) = torch.autograd.grad( + (reference * cotangent).sum(), + reference_edge, + ) + torch.testing.assert_close(output, reference, atol=2e-6, rtol=2e-6) + torch.testing.assert_close( + gradient, + reference_gradient, + atol=8e-6 if channels >= 64 else 3e-6, + rtol=1e-4 if channels >= 64 else 3e-5, + ) + + +@_GPU +@pytest.mark.parametrize("channels", [64, 128]) +def test_wide_backward_is_deterministic(channels: int) -> None: + descriptor = _build_descriptor(channels) + graph, atype = _build_graph(descriptor, canonical=True) + arguments = _arguments(descriptor, graph, atype) + output, state = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec, + *arguments, + ) + previous = torch.are_deterministic_algorithms_enabled() + torch.use_deterministic_algorithms(True) + try: + first = torch.ops.deepmd.dpa4c_graph_compress_backward( + torch.ones_like(output), + state, + graph.edge_vec, + *arguments, + ) + second = torch.ops.deepmd.dpa4c_graph_compress_backward( + torch.ones_like(output), + state, + graph.edge_vec, + *arguments, + ) + finally: + torch.use_deterministic_algorithms(previous) + torch.testing.assert_close(first, second, atol=0.0, rtol=0.0) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 16, 32, 64, 128]) +def test_compressed_matches_uncompressed_descriptor(channels: int) -> None: + descriptor = _build_descriptor(channels) + graph, atype = _build_graph(descriptor, canonical=False) + arguments = _arguments(descriptor, graph, atype) + compressed, _state = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec, + *arguments, + ) + reference, _ = descriptor.call_graph( + graph, + atype, + type_embedding=descriptor.type_embedding.call(), + ) + torch.testing.assert_close(compressed, reference, atol=3e-5, rtol=3e-5) + + cotangent = torch.linspace( + -0.7, + 1.3, + compressed.numel(), + dtype=compressed.dtype, + device=compressed.device, + ).reshape_as(compressed) + compressed_edge = graph.edge_vec.detach().clone().requires_grad_(True) + compressed_value, _state = torch.ops.deepmd.dpa4c_graph_compress( + compressed_edge, + *arguments, + ) + (compressed_gradient,) = torch.autograd.grad( + (compressed_value * cotangent).sum(), + compressed_edge, + ) + reference_edge = graph.edge_vec.detach().clone().requires_grad_(True) + reference_graph = dataclasses.replace(graph, edge_vec=reference_edge) + reference_value, _ = descriptor.call_graph( + reference_graph, + atype, + type_embedding=descriptor.type_embedding.call(), + ) + (reference_gradient,) = torch.autograd.grad( + (reference_value * cotangent).sum(), + reference_edge, + ) + torch.testing.assert_close( + compressed_gradient, + reference_gradient, + atol=1e-4, + rtol=5e-4, + ) + + +@_GPU +def test_output_calibration_matches_uncompressed_descriptor() -> None: + descriptor = _build_descriptor(8) + output_width = descriptor.get_dim_out() + mean = torch.linspace( + -0.4, + 0.6, + output_width, + dtype=torch.float32, + device="cuda", + ) + stddev = torch.linspace( + 0.7, + 1.9, + output_width, + dtype=torch.float32, + device="cuda", + ) + descriptor.set_stat_mean_and_stddev(mean, stddev) + graph, atype = _build_graph(descriptor, canonical=False) + reference, _ = descriptor.call_graph(graph, atype) + actual, _state = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec, + *_arguments(descriptor, graph, atype), + ) + torch.testing.assert_close(actual, reference, atol=3e-5, rtol=3e-5) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 64, 128]) +def test_descriptor_compression_routing_and_serialization( + monkeypatch: pytest.MonkeyPatch, + channels: int, +) -> None: + descriptor = _build_descriptor(channels) + graph, atype = _build_graph(descriptor, canonical=False) + reference, _ = descriptor.call_graph(graph, atype) + descriptor.enable_compression(min_nbor_dist=0.5) + monkeypatch.setenv("DP_CUDA_INFER", "1") + actual, _ = descriptor.call_graph(graph, atype) + restored = DescrptDPA4C.deserialize(descriptor.serialize()).cuda().eval() + restored_output, _ = restored.call_graph(graph, atype) + torch.testing.assert_close(actual, reference, atol=3e-5, rtol=3e-5) + torch.testing.assert_close(restored_output, actual) + + +@_GPU +def test_compression_uses_immutable_type_snapshot( + monkeypatch: pytest.MonkeyPatch, +) -> None: + descriptor = _build_descriptor(8) + graph, atype = _build_graph(descriptor, canonical=False) + descriptor.enable_compression(min_nbor_dist=0.5) + monkeypatch.setenv("DP_CUDA_INFER", "1") + reference, _ = descriptor.call_graph(graph, atype) + alternative = descriptor.type_embedding.call().detach().clone() + 0.25 + actual, _ = descriptor.call_graph( + graph, + atype, + type_embedding=alternative, + ) + torch.testing.assert_close(actual, reference, atol=0.0, rtol=0.0) + + monkeypatch.setenv("DP_CUDA_INFER", "0") + portable, _ = descriptor.call_graph( + graph, + atype, + type_embedding=alternative, + ) + assert not torch.allclose(portable, reference) + + +@_GPU +def test_post_compression_statistics_update_snapshot( + monkeypatch: pytest.MonkeyPatch, +) -> None: + descriptor = _build_descriptor(8) + graph, atype = _build_graph(descriptor, canonical=False) + descriptor.enable_compression(min_nbor_dist=0.5) + output_width = descriptor.get_dim_out() + mean = torch.linspace(-0.2, 0.3, output_width, device="cuda") + stddev = torch.linspace(0.8, 1.6, output_width, device="cuda") + descriptor.set_stat_mean_and_stddev(mean, stddev) + + monkeypatch.setenv("DP_CUDA_INFER", "0") + reference, _ = descriptor.call_graph(graph, atype) + monkeypatch.setenv("DP_CUDA_INFER", "1") + actual, _ = descriptor.call_graph(graph, atype) + torch.testing.assert_close(actual, reference, atol=3e-5, rtol=3e-5) + + +@_GPU +def test_compressed_descriptor_cannot_reenter_training() -> None: + descriptor = _build_descriptor(8) + descriptor.enable_compression(min_nbor_dist=0.5) + with pytest.raises(RuntimeError, match="immutable inference snapshot"): + descriptor.train() + + +@_GPU +def test_compression_rejects_float64_descriptor() -> None: + descriptor = ( + DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=8, + lmax=2, + n_radial=8, + precision="float64", + seed=17, + ) + .cuda() + .eval() + ) + with pytest.raises(ValueError, match="requires descriptor precision `float32`"): + descriptor.enable_compression(min_nbor_dist=0.5) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 16, 32, 64, 128]) +@pytest.mark.parametrize("lmax", [2, 3, 4]) +@pytest.mark.parametrize("radial_modes", [0, 2, 4, 8]) +@pytest.mark.parametrize("canonical", [False, True]) +def test_supported_surface_parity( + channels: int, + lmax: int, + radial_modes: int, + canonical: bool, +) -> None: + """Cover the complete compiled surface against the portable equations. + + Each scalar width owns a distinct lane mapping and each angular degree a + distinct instantiation, so the cross product is the contract the operator + advertises rather than a sample of it. + """ + descriptor = _build_descriptor(channels, lmax, radial_modes) + graph, atype = _build_graph(descriptor, canonical) + arguments = _arguments(descriptor, graph, atype) + ensure_registered() + + edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) + output, _state = torch.ops.deepmd.dpa4c_graph_compress(edge_vec, *arguments) + reference, _ = descriptor.call_graph(graph, atype) + torch.testing.assert_close(output, reference, atol=3e-5, rtol=3e-5) + + cotangent = torch.linspace( + -0.7, + 1.3, + output.numel(), + dtype=output.dtype, + device=output.device, + ).reshape_as(output) + (gradient,) = torch.autograd.grad((output * cotangent).sum(), edge_vec) + reference_edge = graph.edge_vec.detach().clone().requires_grad_(True) + reference_value = _cpu_descriptor(reference_edge, *arguments) + (reference_gradient,) = torch.autograd.grad( + (reference_value * cotangent).sum(), + reference_edge, + ) + torch.testing.assert_close(gradient, reference_gradient, atol=8e-6, rtol=1e-4) + + +@_GPU +@pytest.mark.parametrize( + ("index", "replacement", "message"), + [ + (6, torch.zeros(9, 4, 2, device="cuda"), "PairFiLM"), + (9, torch.zeros(8, 5, 5, device="cuda"), "invalid readout matrix"), + (10, torch.zeros(1, 8, dtype=torch.int32, device="cuda"), "degree triples"), + ], +) +def test_operator_rejects_inconsistent_artifacts( + index: int, + replacement: torch.Tensor, + message: str, +) -> None: + """Device-side shape assumptions are enforced at the operator boundary.""" + descriptor = _build_descriptor(8) + graph, atype = _build_graph(descriptor, canonical=False) + arguments = list(_arguments(descriptor, graph, atype)) + arguments[index] = replacement + with pytest.raises(RuntimeError, match=message): + torch.ops.deepmd.dpa4c_graph_compress(graph.edge_vec, *arguments) + + +@_GPU +def test_operator_rejects_unsupported_mode_rank() -> None: + """An unsupported mode rank would overflow the shared mode cache.""" + descriptor = _build_descriptor(8, radial_modes=2) + graph, atype = _build_graph(descriptor, canonical=False) + arguments = list(_arguments(descriptor, graph, atype)) + table = arguments[5] + channels = descriptor.channels + # Three modes keep the table and cache shapes mutually consistent while + # leaving the rank outside the compiled set. + arguments[5] = torch.zeros( + table.shape[0], + 6 * (channels + 3), + dtype=torch.float32, + device="cuda", + ) + arguments[7] = torch.zeros( + arguments[7].shape[0], + channels, + 3, + dtype=torch.float32, + device="cuda", + ) + with pytest.raises(RuntimeError, match="radial_modes must be"): + torch.ops.deepmd.dpa4c_graph_compress(graph.edge_vec, *arguments) + + +@_GPU +def test_generic_compression_restores_input_dtype( + monkeypatch: pytest.MonkeyPatch, +) -> None: + descriptor = _build_descriptor(8) + graph, atype = _build_graph(descriptor, canonical=False) + graph64 = dataclasses.replace(graph, edge_vec=graph.edge_vec.to(torch.float64)) + descriptor.enable_compression(min_nbor_dist=0.5) + monkeypatch.setenv("DP_CUDA_INFER", "0") + reference, _ = descriptor.call_graph(graph64, atype) + monkeypatch.setenv("DP_CUDA_INFER", "1") + actual, _ = descriptor.call_graph(graph64, atype) + assert actual.dtype == graph64.edge_vec.dtype + torch.testing.assert_close(actual, reference, atol=3e-5, rtol=3e-5) + + +@_GPU +def test_radial_table_accuracy() -> None: + descriptor = _build_descriptor(8) + table, info = build_radial_table(descriptor) + radius = torch.linspace(0.0, descriptor.rcut, 2001, device="cuda") + reference = descriptor.radial_embedding(descriptor.radial_basis(radius[:, None])) + actual = _table_lookup( + table, + radius, + float(info[0]), + float(info[1]), + descriptor.channels, + ) + torch.testing.assert_close(actual, reference, atol=2e-6, rtol=2e-6) + + +@_GPU +def test_radial_table_is_c2_at_internal_knots() -> None: + descriptor = _build_descriptor(8) + table, info = build_radial_table(descriptor) + stride = float(info[0]) + knot = 617 * stride + cotangent = torch.linspace( + -0.7, + 1.3, + descriptor.channels, + device="cuda", + ) + + def derivatives(radius_value: float) -> list[torch.Tensor]: + radius = torch.tensor( + [radius_value], + dtype=torch.float32, + device="cuda", + requires_grad=True, + ) + value = ( + _table_lookup( + table, + radius, + stride, + float(info[1]), + descriptor.channels, + )[0] + * cotangent + ).sum() + first = torch.autograd.grad(value, radius, create_graph=True)[0] + second = torch.autograd.grad(first, radius, create_graph=True)[0] + return [value, first[0], second[0]] + + left = derivatives(knot - 1e-6) + right = derivatives(knot + 1e-6) + tolerances = (3e-6, 1e-4, 2e-3) + for lhs, rhs, atol in zip(left, right, tolerances, strict=True): + torch.testing.assert_close(lhs, rhs, atol=atol, rtol=0.0) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 64, 128]) +def test_compressed_cutoff_matches_removed_topology(channels: int) -> None: + descriptor = _build_descriptor(channels) + edge_index = torch.tensor( + [[1, 0], [0, 1]], + dtype=torch.long, + device="cuda", + ) + radius = torch.tensor(descriptor.rcut, device="cuda") + zero = torch.zeros_like(radius) + edge_vec = torch.stack( + [ + torch.stack([radius, zero, zero]), + torch.stack([-radius, zero, zero]), + ] + ).requires_grad_(True) + graph = NeighborGraph( + n_node=torch.tensor([2], dtype=torch.long, device="cuda"), + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=torch.ones(2, dtype=torch.bool, device="cuda"), + ) + graph = attach_edge_csr(graph, 2, canonicalize=False) + atype = torch.zeros(2, dtype=torch.long, device="cuda") + arguments = _arguments(descriptor, graph, atype) + retained, _state = torch.ops.deepmd.dpa4c_graph_compress( + edge_vec, + *arguments, + ) + (gradient,) = torch.autograd.grad(retained.sum(), edge_vec) + + removed_graph = dataclasses.replace( + graph, + edge_mask=torch.zeros_like(graph.edge_mask), + ) + removed_edge = edge_vec.detach().clone().requires_grad_(True) + removed, _removed_state = torch.ops.deepmd.dpa4c_graph_compress( + removed_edge, + *_arguments(descriptor, removed_graph, atype), + ) + (removed_gradient,) = torch.autograd.grad(removed.sum(), removed_edge) + torch.testing.assert_close(retained, removed, atol=2e-6, rtol=2e-6) + torch.testing.assert_close( + gradient, + torch.zeros_like(gradient), + atol=2e-6, + rtol=0.0, + ) + torch.testing.assert_close( + removed_gradient, + torch.zeros_like(removed_gradient), + atol=0.0, + rtol=0.0, + ) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 64, 128]) +def test_in_row_mask_matches_removed_edge(channels: int) -> None: + descriptor = _build_descriptor(channels) + edge_index = torch.tensor( + [[1, 0], [0, 1]], + dtype=torch.long, + device="cuda", + ) + edge_vec = torch.tensor( + [[1.0, 0.2, -0.1], [-0.7, 0.3, 0.4]], + dtype=torch.float32, + device="cuda", + ) + graph = NeighborGraph( + n_node=torch.tensor([2], dtype=torch.long, device="cuda"), + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=torch.ones(2, dtype=torch.bool, device="cuda"), + ) + graph = attach_edge_csr(graph, 2, canonicalize=False) + atype = torch.zeros(2, dtype=torch.long, device="cuda") + + full, _full_state = torch.ops.deepmd.dpa4c_graph_compress( + edge_vec, + *_arguments(descriptor, graph, atype), + ) + masked_graph = dataclasses.replace( + graph, + edge_mask=torch.tensor([True, False], dtype=torch.bool, device="cuda"), + ) + masked_edge = edge_vec.detach().clone().requires_grad_(True) + masked, _masked_state = torch.ops.deepmd.dpa4c_graph_compress( + masked_edge, + *_arguments(descriptor, masked_graph, atype), + ) + cotangent = torch.linspace( + -0.7, + 1.3, + masked.numel(), + dtype=masked.dtype, + device=masked.device, + ).reshape_as(masked) + (masked_gradient,) = torch.autograd.grad( + (masked * cotangent).sum(), + masked_edge, + ) + + removed_edge = edge_vec[:1].detach().clone().requires_grad_(True) + removed_graph = NeighborGraph( + n_node=graph.n_node, + edge_index=edge_index[:, :1].contiguous(), + edge_vec=removed_edge, + edge_mask=torch.ones(1, dtype=torch.bool, device="cuda"), + ) + removed_graph = attach_edge_csr(removed_graph, 2, canonicalize=False) + removed, _removed_state = torch.ops.deepmd.dpa4c_graph_compress( + removed_edge, + *_arguments(descriptor, removed_graph, atype), + ) + (removed_gradient,) = torch.autograd.grad( + (removed * cotangent).sum(), + removed_edge, + ) + + assert not torch.allclose(full, masked) + torch.testing.assert_close(masked, removed, atol=2e-6, rtol=2e-6) + torch.testing.assert_close( + masked_gradient[:1], + removed_gradient, + atol=3e-6, + rtol=3e-5, + ) + torch.testing.assert_close( + masked_gradient[1], + torch.zeros_like(masked_gradient[1]), + atol=0.0, + rtol=0.0, + ) + assert torch.count_nonzero(masked_gradient[0]).item() > 0 + + +@_GPU +@pytest.mark.parametrize("channels", [8, 32]) +def test_padding_type_edge_has_zero_gradient(channels: int) -> None: + descriptor = _build_descriptor(channels) + edge_vec = torch.tensor( + [[1.0, 0.2, -0.1]], + dtype=torch.float32, + device="cuda", + requires_grad=True, + ) + graph = NeighborGraph( + n_node=torch.tensor([2], dtype=torch.long, device="cuda"), + edge_index=torch.tensor([[1], [0]], dtype=torch.long, device="cuda"), + edge_vec=edge_vec, + edge_mask=torch.ones(1, dtype=torch.bool, device="cuda"), + ) + graph = attach_edge_csr(graph, 2, canonicalize=False) + atype = torch.tensor([0, descriptor.ntypes], dtype=torch.long, device="cuda") + + output, _state = torch.ops.deepmd.dpa4c_graph_compress( + edge_vec, + *_arguments(descriptor, graph, atype), + ) + (gradient,) = torch.autograd.grad(output.sum(), edge_vec) + torch.testing.assert_close( + gradient, + torch.zeros_like(gradient), + atol=0.0, + rtol=0.0, + ) + + +@_GPU +def test_int32_edge_indices() -> None: + descriptor = _build_descriptor(8) + graph, atype = _build_graph(descriptor, canonical=False) + graph32 = dataclasses.replace( + graph, + edge_index=graph.edge_index.to(torch.int32), + destination_order=graph.destination_order.to(torch.int32), + ) + output64, _state64 = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec, + *_arguments(descriptor, graph, atype), + ) + output32, _state32 = torch.ops.deepmd.dpa4c_graph_compress( + graph32.edge_vec, + *_arguments(descriptor, graph32, atype), + ) + torch.testing.assert_close(output32, output64) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 64, 128]) +@pytest.mark.parametrize("index_dtype", [torch.int64, torch.uint32]) +def test_compact_canonical_parity( + channels: int, + index_dtype: torch.dtype, +) -> None: + from deepmd.kernels.cuda.dpa4c.canonical import ( + ensure_registered as ensure_canonical_registered, + ) + + descriptor = _build_descriptor(channels) + graph, atype = _build_graph(descriptor, canonical=True) + arguments = _arguments(descriptor, graph, atype) + ensure_canonical_registered() + generic_output, generic_state = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec, + *arguments, + ) + canonical_arguments = ( + graph.edge_index[0].to(index_dtype), + graph.destination_row_ptr, + atype, + *arguments[5:15], + *arguments[16:], + ) + compact_output, compact_state = torch.ops.deepmd.dpa4c_canonical_compress( + graph.edge_vec, + *canonical_arguments, + ) + torch.testing.assert_close(compact_output, generic_output) + torch.testing.assert_close(compact_state, generic_state) + + cotangent = torch.randn_like(generic_output) + generic_gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( + cotangent, + generic_state, + graph.edge_vec, + *arguments, + ) + compact_gradient = torch.ops.deepmd.dpa4c_canonical_compress_backward( + cotangent, + compact_state, + graph.edge_vec, + *canonical_arguments, + ) + torch.testing.assert_close( + compact_gradient, + generic_gradient, + atol=2e-6, + rtol=2e-6, + ) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 128]) +def test_compact_inplace_backward_reuses_state(channels: int) -> None: + from deepmd.kernels.cuda.dpa4c.canonical import ( + ensure_registered as ensure_canonical_registered, + ) + + descriptor = _build_descriptor(channels) + graph, atype = _build_graph(descriptor, canonical=True, node_count=23) + arguments = _arguments(descriptor, graph, atype) + canonical_arguments = ( + graph.edge_index[0].to(torch.uint32), + graph.destination_row_ptr, + atype, + *arguments[5:15], + *arguments[16:], + ) + ensure_canonical_registered() + output, state = torch.ops.deepmd.dpa4c_canonical_compress( + graph.edge_vec, + *canonical_arguments, + ) + cotangent = torch.randn_like(output) + reference = torch.ops.deepmd.dpa4c_canonical_compress_backward( + cotangent, + state, + graph.edge_vec, + *canonical_arguments, + ) + inplace_state = state.clone() + actual = torch.ops.deepmd.dpa4c_canonical_compress_backward_inplace( + cotangent, + inplace_state, + graph.edge_vec, + *canonical_arguments, + ) + torch.testing.assert_close(actual, reference, atol=0.0, rtol=0.0) + assert not torch.equal(inplace_state, state) + + +@_GPU +@pytest.mark.parametrize("channels", [8, 128]) +@pytest.mark.parametrize("fitting_width", [32, 64, 128, 192, 256]) +@pytest.mark.parametrize("fitting_depth", [1, 2, 3]) +def test_fused_energy_force_parity( + channels: int, + fitting_width: int, + fitting_depth: int, +) -> None: + from deepmd.kernels.cuda.edge_force_virial import ( + edge_force_virial, + ) + from deepmd.pt_expt.fitting.ener_fitting import ( + EnergyFittingNet, + ) + + descriptor = _build_descriptor(channels) + graph, atype = _build_graph(descriptor, canonical=False) + descriptor._set_compression(build_compression_artifacts(descriptor)) + fitting = ( + EnergyFittingNet( + ntypes=2, + dim_descrpt=descriptor.get_dim_out(), + neuron=[fitting_width] * fitting_depth, + resnet_dt=False, + activation_function="silu", + precision="float32", + mixed_types=True, + seed=29, + ) + .cuda() + .eval() + ) + fitting.bias_atom_e = torch.tensor( + [[0.3], [-0.2]], + dtype=torch.float64, + device="cuda", + ) + ownership = torch.ones(atype.shape[0], dtype=torch.bool, device="cuda") + fused = dpa4c_graph_compress_energy_force( + descriptor, + fitting, + graph, + atype, + ownership, + fitting.bias_atom_e[:, 0], + atype.shape[0], + True, + ) + + edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) + arguments = _arguments(descriptor, graph, atype) + node_descriptor, _state = torch.ops.deepmd.dpa4c_graph_compress( + edge_vec, + *arguments, + ) + atom_energy = fitting.call_graph(node_descriptor, atype)[fitting.var_name] + (edge_gradient,) = torch.autograd.grad(atom_energy.sum(), edge_vec) + force, atom_virial, virial = edge_force_virial( + edge_gradient, + edge_vec.detach(), + graph.edge_index, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + graph.n_node, + atype.shape[0], + True, + ) + torch.testing.assert_close( + fused[1], + atom_energy.to(torch.float64), + atol=1e-6, + rtol=1e-6, + ) + torch.testing.assert_close(fused[2], force, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(fused[3], virial, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(fused[4], atom_virial, atol=1e-6, rtol=1e-5) + + +class _ExportModule(torch.nn.Module): + def forward(self, edge_vec: torch.Tensor, *arguments): + return torch.ops.deepmd.dpa4c_graph_compress(edge_vec, *arguments) + + +@_GPU +def test_torch_export() -> None: + descriptor = _build_descriptor(8) + graph, atype = _build_graph(descriptor, canonical=False) + arguments = _arguments(descriptor, graph, atype) + module = _ExportModule().cuda().eval() + exported = torch.export.export( + module, + (graph.edge_vec, *arguments), + strict=False, + ) + actual = exported.module()(graph.edge_vec, *arguments) + reference = module(graph.edge_vec, *arguments) + torch.testing.assert_close(actual, reference) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_fused_energy_force_refuses_ineligible_fitting() -> None: + """The fused path must refuse a network the operator cannot represent. + + A layer timestep has no representation in the fused fitting operator, + which would otherwise evaluate the network without it and return an + energy that no reference path produces. + """ + from deepmd.pt_expt.fitting.ener_fitting import ( + EnergyFittingNet, + ) + + descriptor = _build_descriptor(8) + graph, atype = _build_graph(descriptor, canonical=False) + descriptor._set_compression(build_compression_artifacts(descriptor)) + fitting = ( + EnergyFittingNet( + ntypes=2, + dim_descrpt=descriptor.get_dim_out(), + neuron=[64, 64], + resnet_dt=True, + activation_function="silu", + precision="float32", + mixed_types=True, + seed=29, + ) + .cuda() + .eval() + ) + ownership = torch.ones(atype.shape[0], dtype=torch.bool, device="cuda") + with pytest.raises(ValueError, match="cannot reproduce this network"): + dpa4c_graph_compress_energy_force( + descriptor, + fitting, + graph, + atype, + ownership, + fitting.bias_atom_e[:, 0], + atype.shape[0], + True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("channels", [8, 64]) +def test_compact_canonical_tiling_is_equivalent( + channels: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Node tiling must not change energy, force or virial. + + Every node tile owns a contiguous span of the destination-sorted edge + axis, so the runs partition the work rather than splitting any reduction. + """ + from deepmd.kernels.cuda.dpa4c.canonical import ( + dpa4c_canonical_compress_energy_force, + ) + from deepmd.pt_expt.fitting.ener_fitting import ( + EnergyFittingNet, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + ) + + descriptor = _build_descriptor(channels) + neighbor_graph, atype = _build_graph(descriptor, canonical=True, node_count=97) + graph = canonical_graph_from_neighbor_graph( + dataclasses.replace(neighbor_graph, n_local=neighbor_graph.n_node) + ) + descriptor._set_compression(build_compression_artifacts(descriptor)) + fitting = ( + EnergyFittingNet( + ntypes=2, + dim_descrpt=descriptor.get_dim_out(), + neuron=[64, 64], + resnet_dt=False, + activation_function="silu", + precision="float32", + mixed_types=True, + seed=29, + ) + .cuda() + .eval() + ) + ownership = torch.ones(atype.shape[0], dtype=torch.bool, device="cuda") + + def run() -> tuple[torch.Tensor, ...]: + return dpa4c_canonical_compress_energy_force( + descriptor, + fitting, + graph, + atype, + ownership, + fitting.bias_atom_e[:, 0], + True, + ) + + monkeypatch.setenv("DP_NODE_TILE", "0") + reference = run() + for tile in ("7", "32", "96"): + monkeypatch.setenv("DP_NODE_TILE", tile) + for actual, expected in zip(run(), reference, strict=True): + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-6) diff --git a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py new file mode 100644 index 0000000000..002f870382 --- /dev/null +++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from types import ( + SimpleNamespace, +) + +import pytest +import torch + +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.utils import ( + env, +) +from deepmd.pt_expt.utils.serialization import ( + _resolve_lower_kind, + _trace_and_export, + build_synthetic_graph_inputs, +) + + +def _config() -> dict: + return { + "type_map": ["A", "B"], + "descriptor": { + "type": "dpa4c", + "rcut": 3.0, + "channels": 16, + "lmax": 4, + "n_radial": 4, + "precision": "float64", + "seed": 17, + }, + "fitting_net": { + "type": "ener", + "neuron": [16, 16], + "precision": "float64", + "seed": 19, + }, + } + + +def _compressed_config(channels: int = 8) -> dict: + config = _config() + descriptor = config["descriptor"] + descriptor["channels"] = channels + descriptor["lmax"] = 2 + descriptor["n_radial"] = 8 + descriptor["precision"] = "float32" + fitting = config["fitting_net"] + fitting["neuron"] = [32, 32] + fitting["activation_function"] = "silu" + fitting["precision"] = "float32" + return config + + +def _run_graph(model: torch.nn.Module) -> dict[str, torch.Tensor]: + sample = build_synthetic_graph_inputs( + model, + e_max=None, + nframes=2, + nloc=7, + dtype=torch.float64, + device=env.DEVICE, + ) + ( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + ) = sample + return model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=True, + do_atomic_virial=True, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + + +def test_graph_lower_energy_force_are_finite() -> None: + model = get_model(_config()).to(env.DEVICE).eval() + assert model.get_descriptor().uses_graph_lower() + result = _run_graph(model) + for key in ( + "energy", + "energy_redu", + "energy_derv_r", + "energy_derv_c_redu", + ): + assert key in result + assert torch.isfinite(result[key]).all(), key + + +def test_graph_force_loss_trains_descriptor() -> None: + model = get_model(_config()).to(env.DEVICE).train() + result = _run_graph(model) + loss = result["energy_redu"].square().mean() + loss = loss + result["energy_derv_r"].square().mean() + loss.backward() + descriptor = model.get_descriptor() + gradients = { + name: parameter.grad for name, parameter in descriptor.named_parameters() + } + for name, gradient in gradients.items(): + assert gradient is not None, name + assert torch.isfinite(gradient).all(), name + + +def test_graph_export() -> None: + model = get_model(_config()).to("cpu").eval() + exported, _metadata, _model_json, _output_keys = _trace_and_export( + {"model": model.serialize()}, + lower_kind="graph", + do_atomic_virial=True, + ) + assert isinstance(exported, torch.export.ExportedProgram) + + +def test_compressed_graph_export(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DP_CUDA_INFER", "1") + model = get_model(_compressed_config()).to("cpu").eval() + model.get_descriptor().enable_compression(min_nbor_dist=0.5) + exported, metadata, _model_json, _output_keys = _trace_and_export( + {"model": model.serialize()}, + lower_kind="graph", + do_atomic_virial=True, + ) + assert isinstance(exported, torch.export.ExportedProgram) + assert metadata["graph_edge_dtype"] == "float32" + + +@pytest.mark.parametrize("channels", [8, 64]) +def test_compact_canonical_graph_export( + monkeypatch: pytest.MonkeyPatch, + channels: int, +) -> None: + monkeypatch.setenv("DP_CUDA_INFER", "2") + model = get_model(_compressed_config(channels)).to("cpu").eval() + model.get_descriptor().enable_compression(min_nbor_dist=0.5) + exported, metadata, _model_json, _output_keys = _trace_and_export( + {"model": model.serialize()}, + lower_kind="dpa4c_canonical", + do_atomic_virial=True, + ) + assert isinstance(exported, torch.export.ExportedProgram) + assert metadata["lower_input_kind"] == "dpa4c_canonical" + assert metadata["graph_edge_dtype"] == "float32" + assert metadata["canonical_index_dtype"] == "uint32" + + +@pytest.mark.parametrize("channels", [8, 64, 128]) +def test_auto_lower_kind_selects_compact_canonical(channels: int) -> None: + model = get_model(_compressed_config(channels)).to("cpu").eval() + model.get_descriptor().enable_compression(min_nbor_dist=0.5) + data = {"model": model.serialize()} + assert _resolve_lower_kind("model.pt2", data, "auto") == "dpa4c_canonical" + + +def test_compact_canonical_eligibility_rejects_other_descriptors() -> None: + from deepmd.kernels.cuda.dpa4c.canonical import ( + canonical_model_eligible, + ) + + model = SimpleNamespace( + atomic_model=SimpleNamespace( + descriptor=SimpleNamespace(compress=True), + fitting_net=object(), + ) + ) + assert not canonical_model_eligible(model) + + +@pytest.mark.parametrize("channels", [8, 64]) +def test_compressed_level_two_matches_autograd( + monkeypatch: pytest.MonkeyPatch, + channels: int, +) -> None: + model = get_model(_compressed_config(channels)).to(env.DEVICE).eval() + model.get_descriptor().enable_compression(min_nbor_dist=0.5) + monkeypatch.setenv("DP_CUDA_INFER", "1") + reference = _run_graph(model) + monkeypatch.setenv("DP_CUDA_INFER", "2") + actual = _run_graph(model) + for key in ( + "energy", + "energy_redu", + "energy_derv_r", + "energy_derv_c", + "energy_derv_c_redu", + ): + torch.testing.assert_close(actual[key], reference[key]) diff --git a/source/tests/pt_expt/utils/test_canonical_graph.py b/source/tests/pt_expt/utils/test_canonical_graph.py index e7041d7f6e..3128db789a 100644 --- a/source/tests/pt_expt/utils/test_canonical_graph.py +++ b/source/tests/pt_expt/utils/test_canonical_graph.py @@ -6,6 +6,8 @@ NeighborGraph, ) from deepmd.pt_expt.utils.canonical_graph import ( + UINT32_MAX, + CanonicalGraph, canonical_graph_from_neighbor_graph, validate_canonical_graph_shapes, ) @@ -42,8 +44,32 @@ def test_storage_guards_remain_outside_csr(physical_edge_count: int) -> None: assert compact.source.shape == (2,) assert compact.edge_vec.shape == (2, 3) assert compact.source_order.shape == (2,) - assert compact.source.dtype == torch.int64 - assert compact.source_order.dtype == torch.int64 + assert compact.source.dtype == torch.uint32 + assert compact.source_order.dtype == torch.uint32 assert int(compact.destination_row_ptr[-1]) == physical_edge_count assert int(compact.source_row_ptr[-1]) == physical_edge_count validate_canonical_graph_shapes(compact, 1) + + +def test_storage_exceeding_uint32_range_is_rejected() -> None: + storage_count = UINT32_MAX + 1 + graph = CanonicalGraph( + n_node=torch.empty(1, dtype=torch.int64, device="meta"), + n_local=torch.empty(1, dtype=torch.int64, device="meta"), + source=torch.empty(storage_count, dtype=torch.uint32, device="meta"), + edge_vec=torch.empty( + storage_count, + 3, + dtype=torch.float32, + device="meta", + ), + destination_row_ptr=torch.empty(2, dtype=torch.int64, device="meta"), + source_row_ptr=torch.empty(2, dtype=torch.int64, device="meta"), + source_order=torch.empty( + storage_count, + dtype=torch.uint32, + device="meta", + ), + ) + with pytest.raises(ValueError, match="exceeds the uint32 range"): + validate_canonical_graph_shapes(graph, 1) From 880baf0c752db0976f68500a4e8a6d30a3067a13 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 29 Jul 2026 15:00:08 +0800 Subject: [PATCH 02/10] feat(dpa4c-spin): add end-to-end native-spin support Add per-atom native-spin conditioning to DPA4C from descriptor training through compressed deployment. - implement spin-aware invariant channels, statistics, serialization, evaluation, and validation - expose magnetic outputs through Python, C/C++, and LAMMPS/Kokkos interfaces - extend the compressed CUDA path and fused reductions for magnetic forces - document the model contract and provide a non-spin-dynamics LAMMPS example --- .../dpmodel/atomic_model/dp_atomic_model.py | 4 + .../atomic_model/linear_atomic_model.py | 9 + .../atomic_model/make_base_atomic_model.py | 11 + .../atomic_model/pairtab_atomic_model.py | 11 + deepmd/dpmodel/descriptor/dpa4c.py | 536 +++++++-- .../dpmodel/descriptor/dpa4c_nn/__init__.py | 8 + .../dpmodel/descriptor/dpa4c_nn/pair_film.py | 132 ++- deepmd/dpmodel/descriptor/dpa4c_nn/spin.py | 790 +++++++++++++ .../descriptor/make_base_descriptor.py | 15 + deepmd/dpmodel/model/make_model.py | 4 + deepmd/infer/model_test/ener.py | 8 +- deepmd/kernels/cuda/dpa1/canonical.py | 3 +- deepmd/kernels/cuda/dpa1/graph_compress.py | 3 +- .../kernels/cuda/dpa1/graph_energy_force.py | 3 +- deepmd/kernels/cuda/dpa4c/canonical.py | 132 ++- deepmd/kernels/cuda/dpa4c/graph_compress.py | 622 +++++++++- deepmd/kernels/cuda/edge_force_virial.py | 48 +- deepmd/pt_expt/descriptor/dpa1.py | 64 +- deepmd/pt_expt/descriptor/dpa4c.py | 66 +- deepmd/pt_expt/entrypoints/compress.py | 10 +- deepmd/pt_expt/infer/deep_eval.py | 82 +- deepmd/pt_expt/model/edge_transform_output.py | 3 +- deepmd/pt_expt/model/ener_model.py | 10 +- deepmd/pt_expt/model/make_model.py | 23 +- deepmd/pt_expt/model/native_spin_model.py | 88 ++ deepmd/pt_expt/train/training.py | 20 +- deepmd/pt_expt/train/validation.py | 8 +- deepmd/pt_expt/utils/serialization.py | 97 +- deepmd/utils/argcheck.py | 9 +- deepmd/utils/eval_metrics.py | 93 +- doc/model/dpa4c.md | 130 ++ examples/spin/dpa4c/input.json | 86 ++ examples/spin/dpa4c/lmp/README.md | 65 + examples/spin/dpa4c/lmp/in.lammps | 66 ++ examples/spin/dpa4c/lmp/init.data | 80 ++ source/api_c/include/c_api.h | 53 + source/api_c/include/deepmd.hpp | 45 + source/api_c/src/c_api.cc | 40 + source/api_cc/CMakeLists.txt | 6 +- source/api_cc/include/DeepSpin.h | 91 ++ source/api_cc/include/NativeSpinPTExpt.h | 370 ++++++ source/api_cc/include/commonPT.h | 19 + source/api_cc/src/DeepPotPTExpt.cc | 6 + source/api_cc/src/DeepPotPTExptPlugin.cc | 57 + source/api_cc/src/DeepSpin.cc | 71 ++ source/api_cc/src/NativeSpinPTExpt.cc | 1051 +++++++++++++++++ source/lmp/compact_canonical_graph_kokkos.h | 477 ++++++++ source/lmp/pair_deepmd_kokkos.cpp | 444 +------ source/lmp/pair_deepmd_kokkos.h | 54 +- source/lmp/pair_dpa4spin.cpp | 517 ++++++++ source/lmp/pair_dpa4spin.h | 106 ++ source/lmp/pair_dpa4spin_kokkos.cpp | 475 ++++++++ source/lmp/pair_dpa4spin_kokkos.h | 109 ++ source/op/pt/dpa1_graph_energy_force.cu | 16 +- source/op/pt/dpa4c_graph_compress.cu | 543 ++++++--- source/op/pt/dpa4c_graph_compress.cuh | 116 +- source/op/pt/dpa4c_graph_compress_kernel.cuh | 703 ++++++++++- source/op/pt/dpa4c_graph_compress_launch.h | 121 +- source/op/pt/edge_force_virial.cu | 165 ++- source/op/pt/graph_ops.h | 31 +- .../common/dpmodel/test_descriptor_dpa4c.py | 611 +++++++++- source/tests/common/test_examples.py | 1 + .../pt_expt/descriptor/test_dpa1_cuda.py | 92 +- source/tests/pt_expt/descriptor/test_dpa4c.py | 150 ++- .../pt_expt/descriptor/test_dpa4c_cuda.py | 313 ++++- .../pt_expt/model/test_dpa4_native_spin.py | 4 +- .../pt_expt/model/test_dpa4c_graph_lower.py | 121 +- .../tests/pt_expt/model/test_zbl_bridging.py | 4 +- 68 files changed, 9201 insertions(+), 1120 deletions(-) create mode 100644 deepmd/dpmodel/descriptor/dpa4c_nn/spin.py create mode 100644 examples/spin/dpa4c/input.json create mode 100644 examples/spin/dpa4c/lmp/README.md create mode 100644 examples/spin/dpa4c/lmp/in.lammps create mode 100644 examples/spin/dpa4c/lmp/init.data create mode 100644 source/api_cc/include/NativeSpinPTExpt.h create mode 100644 source/api_cc/src/NativeSpinPTExpt.cc create mode 100644 source/lmp/compact_canonical_graph_kokkos.h create mode 100644 source/lmp/pair_dpa4spin.cpp create mode 100644 source/lmp/pair_dpa4spin.h create mode 100644 source/lmp/pair_dpa4spin_kokkos.cpp create mode 100644 source/lmp/pair_dpa4spin_kokkos.h diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 0406eb4b91..dbafd00b0a 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -182,6 +182,10 @@ def supports_graph_export(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.supports_graph_export()) + def compression_needs_min_nbor_dist(self) -> bool: + """Delegates to this model's own descriptor.""" + return bool(self.descriptor.compression_needs_min_nbor_dist()) + def supports_native_spin(self) -> bool: """Delegates to this model's own descriptor (cached at construction).""" return self._supports_native_spin diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index bb42ec37ae..6d9a93ac65 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -303,6 +303,15 @@ def enable_compression( check_frequency, ) + def compression_needs_min_nbor_dist(self) -> bool: + """Required as soon as ANY child consumes it. + + The statistic is measured once and handed to every child, so a single + child that tabulates from the shortest observed distance keeps the + neighbor-statistics pass for the whole composition. + """ + return any(m.compression_needs_min_nbor_dist() for m in self.models) + def uses_graph_lower(self) -> bool: """Graph-capable iff EVERY child supports the graph lower. diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index f5deadde11..2d4032f09e 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -190,6 +190,17 @@ def enable_compression( """ raise NotImplementedError("This atomi model doesn't support compression!") + def compression_needs_min_nbor_dist(self) -> bool: + """Whether :meth:`enable_compression` consumes ``min_nbor_dist``. + + Returns + ------- + bool + Concrete default ``True``, so a model that does not report + otherwise keeps the neighbor-statistics pass. + """ + return True + def make_atom_mask( self, atype: t_tensor, diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index cea0403812..8a466b9702 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -505,3 +505,14 @@ def enable_compression( ) -> None: """Pairtab model does not support compression.""" pass + + def compression_needs_min_nbor_dist(self) -> bool: + """Return whether compression consumes the minimum neighbor distance. + + Returns + ------- + bool + Always ``False``. The tabulated pair potential carries its own + domain, so compression is a no-op here. + """ + return False diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index 4f13407ef8..f854259c9f 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -75,10 +75,13 @@ from .dpa4c_nn import ( InvariantReadout, OrderedPairFiLM, + SpinChannels, build_angular_basis, build_moment_indices, + degree_offsets, derive_bispectrum_ranks, derive_degree_channels, + derive_spin_channels, ) if TYPE_CHECKING: @@ -166,8 +169,14 @@ class DescrptDPA4C(NativeOP, BaseDescriptor): Atom-type names. seed Random seed. + use_spin + Per-type flags marking which atom types carry a magnetic moment. When + given, the descriptor conditions on a per-node spin vector and + declares :meth:`supports_native_spin`. ``None`` reproduces the + spin-free descriptor exactly. spin Reserved for descriptor API compatibility; only ``None`` is supported. + Native spin is configured through ``use_spin``. Raises ------ @@ -197,6 +206,8 @@ class DescrptDPA4C(NativeOP, BaseDescriptor): "info", "pair_film", "pair_mixing", + "spin_pair", + "spin_type", "type_embedding", "readout_matrices", "coupling_meta", @@ -221,11 +232,15 @@ def __init__( trainable: bool = True, type_map: list[str] | None = None, seed: int | list[int] | None = None, + use_spin: list[bool] | None = None, spin: None = None, ) -> None: # === Step 1. Validate the public architecture contract === if spin is not None: - raise NotImplementedError("DPA4C does not support spin inputs.") + raise NotImplementedError( + "DPA4C configures native spin through `use_spin`; the `spin` " + "argument of the common descriptor ABI is not supported." + ) if rcut <= 0.0: raise ValueError(f"`rcut` must be positive, got {rcut}") if ntypes <= 0: @@ -258,6 +273,12 @@ def __init__( self.trainable = bool(trainable) self.type_map = type_map self.seed = seed + self.use_spin = None if use_spin is None else [bool(flag) for flag in use_spin] + # The spin branch reads the leading channels of the shared radial map, + # so its width is derived rather than exposed. + self.spin_channels = ( + 0 if self.use_spin is None else derive_spin_channels(degree_channels) + ) radial_hidden = resolve_swiglu_hidden_width(self.channels) # === Step 3. Build the shared DPA4 edge representation === @@ -310,6 +331,7 @@ def __init__( self.pair_film = OrderedPairFiLM( self.channels, radial_modes=self.radial_modes, + spin_channels=self.spin_channels, precision=self.precision, trainable=self.trainable, seed=child_seed(seed, 3), @@ -321,14 +343,28 @@ def __init__( trainable=self.trainable, seed=child_seed(seed, 4), ) + self.spin = ( + None + if self.use_spin is None + else SpinChannels( + self.ntypes, + self.degree_channels, + self.use_spin, + precision=self.precision, + trainable=self.trainable, + seed=child_seed(seed, 5), + ) + ) # === Step 4. Lay out the flat moment payload === # Degree zero owns the leading `channels` entries of the flat layout, # so the non-scalar block is exactly its complement and needs no - # separate degree index. + # separate degree index. The spin families, when present, are appended + # after the geometric degrees, leaving every geometric offset intact. channel_index, harmonic_index = build_moment_indices(self.degree_channels) self.angular_channel_index = channel_index[self.channels :] self.angular_harmonic_index = harmonic_index[self.channels :] + self.degree_offsets = degree_offsets(self.degree_channels) # === Step 5. Initialize the output calibration state === mean = np.zeros( @@ -348,6 +384,7 @@ def call_graph( atype: Array, type_embedding: Array | None = None, comm_dict: dict | None = None, + spin: Array | None = None, ) -> tuple[Array, None]: """Evaluate DPA4C on a flat neighbor graph. @@ -368,6 +405,11 @@ def call_graph( Communication metadata accepted by the common graph ABI. DPA4C does not read source-node features, so no halo-feature exchange is required and this argument is unused. + spin + Per-node spin vectors with shape ``(N, 3)`` on the same flat node + axis as ``atype``, including ghost and padding rows. Mandatory + when the descriptor is configured with ``use_spin`` and ignored + otherwise. Returns ------- @@ -376,6 +418,11 @@ def call_graph( ``(N, get_dim_out())`` and the same floating dtype as ``edge_vec``. rot_mat ``None``. DPA4C does not expose an equivariant fitting input. + + Raises + ------ + ValueError + If the descriptor is spin conditioned and ``spin`` is absent. """ del comm_dict # === Step 1. Resolve type features and compute precision === @@ -391,13 +438,48 @@ def call_graph( ) # === Step 2. Evaluate the graph-native equations === - descriptor, _ = self.evaluate_graph(graph, atype, type_embedding) + descriptor, _ = self.evaluate_graph(graph, atype, type_embedding, spin) # === Step 3. Restore the graph input dtype === if descriptor.dtype != in_dtype: descriptor = xp.astype(descriptor, in_dtype) return descriptor, None + def require_spin(self, spin: Array | None) -> Array: + """Return the per-node moment a spin-conditioned descriptor must receive. + + A missing moment is an error rather than a vanishing one. Substituting + zeros would report an identically zero magnetic force, which in + molecular dynamics is indistinguishable from frozen moments under a + plausible energy. A corpus that is only partially labelled is admitted + through ``model.spin.allow_missing_label``, which relaxes the ``spin`` + data requirement to optional with a zero default, so the data pipeline + supplies an explicit zero moment and this contract still holds. + + Parameters + ---------- + spin + Per-node spin vectors with shape ``(N, 3)``, or ``None``. + + Returns + ------- + Array + The same moments, unchanged. + + Raises + ------ + ValueError + If ``spin`` is ``None``. + """ + if spin is None: + raise ValueError( + "A spin-conditioned DPA4C requires a per-node magnetic " + "moment. Set `model.spin.allow_missing_label` to admit " + "systems that carry no spin label; the data pipeline then " + "supplies an explicit zero moment." + ) + return spin + @cast_precision def call( self, @@ -489,6 +571,7 @@ def evaluate_graph( graph: Any, atype: Array, type_embedding: Array, + spin: Array | None = None, ) -> tuple[Array, Array]: """Evaluate the graph-native descriptor equations. @@ -503,6 +586,9 @@ def evaluate_graph( Flat node types with shape ``(N,)``. type_embedding Complete type table with shape ``(ntypes + 1, channels)``. + spin + Per-node spin vectors with shape ``(N, 3)``, mandatory for a + spin-conditioned descriptor and ignored otherwise. Returns ------- @@ -510,6 +596,11 @@ def evaluate_graph( Invariant node features with shape ``(N, get_dim_out())``. envelope Masked per-edge C³ envelope with shape ``(E, 1)``. + + Raises + ------ + ValueError + If the descriptor is spin conditioned and ``spin`` is absent. """ xp = array_api_compat.array_namespace(graph.edge_vec) @@ -529,23 +620,38 @@ def evaluate_graph( center_type_embedding = self.gather_rows(type_embedding, atype, xp) pair_tables = self.pair_film.call(type_embedding) - # === Step 2. Build the masked edge amplitudes and harmonics === - amplitude, basis, envelope = self.build_edge_features( + # === Step 2. Condition the per-node spin === + # The mask and the reference magnitude are applied once, so every + # downstream spin route inherits them and the magnetic force of a + # non-magnetic type vanishes identically rather than numerically. + conditioned_spin = ( + None + if self.spin is None + else self.spin.conditioned_spin(self.require_spin(spin), atype) + ) + + # === Step 3. Build the masked edge amplitudes and harmonics === + amplitude, basis, envelope, spin_payload = self.build_edge_features( graph, atype, - *pair_tables, + pair_tables, + conditioned_spin, ) - # === Step 3. Reduce the degree-wise moments === + # === Step 4. Reduce the degree-wise moments === moments, divisors = self.aggregate_moments( amplitude, basis, envelope, + spin_payload, + None + if conditioned_spin is None + else self.spin.onsite_payload(conditioned_spin, atype), dst, n_total, ) - # === Step 4. Build calibrated invariant features === + # === Step 5. Build calibrated invariant features === return ( self.build_invariant_descriptor( moments, @@ -559,11 +665,10 @@ def build_edge_features( self, graph: Any, atype: Array, - pair_scale: Array, - pair_shift: Array, - pair_mixing: Array | None, - ) -> tuple[Array, Array, Array]: - r"""Build the enveloped edge amplitudes and the masked harmonics. + pair_tables: tuple, + conditioned_spin: Array | None = None, + ) -> tuple[Array, Array, Array, Array | None]: + r"""Build the enveloped edge amplitudes, harmonics, and spin payload. The ordered type pair :math:`(a,b)` rescales the one shared radial function :math:`g` and mixes the :math:`R` shared mode profiles @@ -577,22 +682,27 @@ def build_edge_features( +\sum_\mu U_{ab,c\mu}q_\mu(\rho_{ij}) \Bigr). + The spin payload reuses the radial map and the ordered pair index of + this same stage, so the whole per-edge computation reads the radial + table exactly once. + Parameters ---------- graph Neighbor graph in descriptor compute precision. atype Flat node types with shape ``(N,)``. - pair_scale - Ordered radial scales with shape - ``((ntypes + 1) ** 2, channels)``. - pair_shift - Ordered radial shifts with shape - ``((ntypes + 1) ** 2, channels)``. - pair_mixing - Ordered mode-mixing table with shape - ``((ntypes + 1) ** 2, channels, radial_modes)``, or ``None`` when - ``radial_modes`` is zero. + pair_tables + Ordered pair cache produced by + :meth:`~deepmd.dpmodel.descriptor.dpa4c_nn.pair_film.OrderedPairFiLM.call`: + radial scale and shift with shape + ``((ntypes + 1) ** 2, channels)``, the mode-mixing table with + shape ``((ntypes + 1) ** 2, channels, radial_modes)`` or ``None``, + and the ordered spin scale and shift with shape + ``((ntypes + 1) ** 2, spin_channels)`` or ``None``. + conditioned_spin + Conditioned per-node spin with shape ``(N, 3)``, or + ``None`` for a spin-free descriptor. Returns ------- @@ -602,11 +712,16 @@ def build_edge_features( Masked Cartesian harmonics with shape ``(E, (lmax + 1) ** 2)``. envelope Masked C³ envelope with shape ``(E,)``. + spin_payload + Masked per-edge spin payload with shape + ``(E, spin.edge_width)``, or ``None`` for a spin-free descriptor. """ from deepmd.dpmodel.utils.neighbor_graph import ( apply_pair_exclusion, ) + pair_scale, pair_shift, pair_mixing, spin_scale, spin_shift = pair_tables + # === Step 1. Merge graph and descriptor-level exclusion masks === graph = apply_pair_exclusion(graph, atype, self.emask) xp = array_api_compat.array_namespace(graph.edge_vec) @@ -653,11 +768,32 @@ def build_edge_features( modes = self.radial_mode_head(radial_hidden) # (E, R) amplitude = amplitude + xp.sum(mixing * modes[:, None, :], axis=-1) - # === Step 5. Gate the amplitude and build the masked harmonics === + # === Step 5. Build the spin payload on the same radial evaluation === + # The bond-projected family reads the same regularized direction as the + # harmonics, so the spin branch contributes an angular term to the + # coordinate gradient alongside the radial one. + spin_payload = ( + None + if conditioned_spin is None + else self.spin.edge_payload( + conditioned_spin, + atype, + src, + direction, + radial, + envelope[:, 0], + spin_scale, + spin_shift, + pair_index, + ) + ) + + # === Step 6. Gate the amplitude and build the masked harmonics === return ( amplitude * envelope, self.build_angular_basis(direction) * mask, envelope[:, 0], + spin_payload, ) def aggregate_moments( @@ -665,6 +801,8 @@ def aggregate_moments( amplitude: Array, basis: Array, envelope: Array, + spin_payload: Array | None, + spin_onsite: Array | None, dst: Array, n_total: int, ) -> tuple[Array, Array]: @@ -694,6 +832,16 @@ def aggregate_moments( Masked Cartesian harmonics with shape ``(E, (lmax + 1) ** 2)``. envelope Masked C³ envelope with shape ``(E,)``. + spin_payload + Masked per-edge spin payload with shape ``(E, spin.edge_width)``, + or ``None``. It carries the same squared envelope as every + non-scalar geometric moment and therefore shares the normalizer + :math:`n^{(+)}`. + spin_onsite + Node-local on-site spin payload with shape + ``(N, spin.node_width)``, or ``None``. It is appended after the + division so that the invariants it enters carry exactly one + neighborhood normalizer, contributed by its neighbour partner. dst Destination node indices with shape ``(E,)``. n_total @@ -703,7 +851,8 @@ def aggregate_moments( ------- moments Flat normalized moments with shape ``(N, S)``, where - ``S = sum((2 * l + 1) * degree_channels[l])``. + ``S = sum((2 * l + 1) * degree_channels[l])`` plus the spin + moment width when the descriptor is spin conditioned. divisors The two divisors :math:`1/n^{(0)}` and :math:`1/n^{(+)}` with shape ``(N, 2)``. They are retained because normalization is otherwise @@ -733,32 +882,28 @@ def aggregate_moments( # single envelope, while the non-scalar block gathers an amplitude and # a harmonic per flat moment coordinate and carries a second envelope. envelope_squared = envelope * envelope - payload = xp.concat( - [ - envelope_squared[:, None], - (envelope_squared * envelope_squared)[:, None], - amplitude, - xp.take(amplitude, channel_index, axis=1) - * xp.take(basis, harmonic_index, axis=1) - * envelope[:, None], - ], - axis=1, - ) - reduced = segment_sum(payload, dst, n_total) + parts = [ + envelope_squared[:, None], + (envelope_squared * envelope_squared)[:, None], + amplitude, + xp.take(amplitude, channel_index, axis=1) + * xp.take(basis, harmonic_index, axis=1) + * envelope[:, None], + ] + if spin_payload is not None: + parts.append(spin_payload) + reduced = segment_sum(xp.concat(parts, axis=1), dst, n_total) scalar_end = 2 + self.channels floor = self._DEGREE_NORM_FLOOR divisors = xp.sqrt(reduced[:, :2] + floor) - return ( - xp.concat( - [ - reduced[:, 2:scalar_end] / divisors[:, :1], - reduced[:, scalar_end:] / divisors[:, 1:], - ], - axis=1, - ), - divisors, - ) + normalized = [ + reduced[:, 2:scalar_end] / divisors[:, :1], + reduced[:, scalar_end:] / divisors[:, 1:], + ] + if spin_onsite is not None: + normalized.append(spin_onsite) + return xp.concat(normalized, axis=1), divisors def build_invariant_descriptor( self, @@ -789,10 +934,22 @@ def build_invariant_descriptor( """ xp = array_api_compat.array_namespace(moments) device = array_api_compat.device(moments) - descriptor = xp.concat( - [self.readout.call(moments), divisors, center_type_embedding], - axis=-1, - ) + blocks = [self.readout.call(moments)] + if self.spin is not None: + geometric_width = self.degree_offsets[-1] + blocks.append( + self.spin.call( + moments[:, geometric_width:], + xp.reshape( + moments[:, self.degree_offsets[2] : self.degree_offsets[3]], + (moments.shape[0], 5, self.degree_channels[2]), + ), + ) + ) + # The two divisors close the geometric block, so the spin invariants + # precede them and the center-type tail keeps its trailing position. + blocks.extend([divisors, center_type_embedding]) + descriptor = xp.concat(blocks, axis=-1) mean = xp_asarray_nodetach(xp, self.mean, device=device) stddev = xp_asarray_nodetach(xp, self.stddev, device=device) return (descriptor - mean[None, :]) / stddev[None, :] @@ -900,6 +1057,7 @@ def share_params( "radial_mode_head", "pair_film", "readout", + "spin", ): setattr(self, name, getattr(base_class, name)) self.mean = base_class.mean @@ -915,13 +1073,14 @@ def structure_signature(self) -> tuple: of a shared module. ``rcut``, ``basis_type``, and ``n_radial`` define the radial basis; ``ntypes`` defines the type table and the ordered pair index space; ``channels``, ``lmax``, and ``radial_modes`` define - every remaining width; ``use_amp`` selects the precision policy that a + every remaining width; ``use_amp`` selects the precision policy that a backend attaches to the shared layers, so a replica that autocasts against layers configured without it would silently lose the effect; ``trainable`` decides whether those layers carry gradients at all; ``type_map`` fixes what the rows - of the shared type table mean. Precision itself enters through its - resolved dtype so that equivalent spellings agree. + of the shared type table mean; ``use_spin`` fixes both the presence + and the row meaning of the shared spin tables. Precision itself enters + through its resolved dtype so that equivalent spellings agree. Branch-local state is deliberately absent. ``exclude_types`` is the only such field: it configures the pair-exclusion mask, which each @@ -943,6 +1102,7 @@ def structure_signature(self) -> tuple: self.use_amp, self.trainable, None if self.type_map is None else tuple(self.type_map), + None if self.use_spin is None else tuple(self.use_spin), np.dtype(PRECISION_DICT[self.precision]).name, ) @@ -1012,10 +1172,19 @@ def compute_input_stats( Parameters ---------- merged - Sampled training systems or a callable returning them. + Sampled training systems or a callable returning them. A + spin-conditioned descriptor additionally requires the per-atom + moment on every system, under either ``model_spin`` or ``spin``. path Optional statistics path. Model-dependent calibration is always recomputed and therefore does not consume this path. + + Raises + ------ + ValueError + If a geometric coordinate is degenerate over the sample, or if a + spin-conditioned descriptor is calibrated on a system that carries + no moment. """ from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph, @@ -1026,6 +1195,12 @@ def compute_input_stats( if not sampled: return + # The reference magnitudes rescale the spin before it reaches the + # descriptor, so they have to be fixed before the output coordinates + # are measured. + if self.spin is not None: + self.spin.set_spin_reference(self._measure_spin_reference(sampled)) + xp = array_api_compat.array_namespace(self.stddev) device = array_api_compat.device(self.stddev) dtype = self.stddev.dtype @@ -1035,48 +1210,35 @@ def compute_input_stats( geometry_dim = self.get_dim_out() - self.channels square_sum = np.zeros(geometry_dim, dtype=np.float64) value_sum = np.zeros(geometry_dim, dtype=np.float64) - value_count = 0 + # A spin coordinate is exactly zero on every node that carries no + # magnetic information, so pooling it over all nodes would scale the + # preconditioner with the magnetic fraction of the sample, that is + # with the stoichiometry rather than with the physics. The spin block + # is therefore averaged over the nodes on which it is active. + # + # Activity is a property of the coordinate block, not of the value. A + # geometric coordinate is legitimately zero on an atom with no + # neighbour inside the cutoff, and counting only nonzero values would + # scale the geometric preconditioner with the vacuum fraction of the + # sample -- the same bias, transposed. The geometric block therefore + # keeps the plain node count. + spin_block = slice(self.readout.get_dim_out(), geometry_dim - 2) + active_count = np.zeros(geometry_dim, dtype=np.float64) try: for system in sampled: - coord_np = to_numpy_array(system["coord"]) - nframes = coord_np.shape[0] - coord_np = np.reshape(coord_np, (nframes, -1, 3)) - atype_np = np.reshape( - to_numpy_array(system["atype"]), - (nframes, -1), - ) - box_value = system.get("box", None) - box_np = ( - None - if box_value is None - else np.reshape(to_numpy_array(box_value), (nframes, -1)) - ) - nstat_frames = min(nframes, self._STAT_FRAMES_PER_SAMPLE) - frame_indices = np.linspace( - 0, - nframes - 1, - num=nstat_frames, - dtype=np.int64, - ) - for frame_index in frame_indices: - coord = xp.asarray( - coord_np[frame_index : frame_index + 1], - dtype=dtype, - device=device, - ) - atype = xp.asarray( - atype_np[frame_index : frame_index + 1], - device=device, - ) + for frame in self._calibration_frames(system): + coord = xp.asarray(frame["coord"], dtype=dtype, device=device) + atype = xp.asarray(frame["atype"], device=device) box = ( None - if box_np is None - else xp.asarray( - box_np[frame_index : frame_index + 1], - dtype=dtype, - device=device, - ) + if frame["box"] is None + else xp.asarray(frame["box"], dtype=dtype, device=device) + ) + spin = ( + None + if frame["spin"] is None + else xp.asarray(frame["spin"], dtype=dtype, device=device) ) graph = build_neighbor_graph( coord, @@ -1084,37 +1246,50 @@ def compute_input_stats( box, self.get_rcut(), ) - output, _ = self.call_graph(graph, xp.reshape(atype, (-1,))) + output, _ = self.call_graph( + graph, + xp.reshape(atype, (-1,)), + spin=None if spin is None else xp.reshape(spin, (-1, 3)), + ) output_np = to_numpy_array(output).reshape( -1, self.get_dim_out(), ) if output_np.shape[0] == 0: continue - square_sum += np.sum( - np.square( - output_np[:, :geometry_dim], - dtype=np.float64, - ), - axis=0, - dtype=np.float64, - ) - value_sum += np.sum( - output_np[:, :geometry_dim], - axis=0, - dtype=np.float64, + geometry = output_np[:, :geometry_dim] + square_sum += np.sum(np.square(geometry, dtype=np.float64), axis=0) + value_sum += np.sum(geometry, axis=0, dtype=np.float64) + active_count += geometry.shape[0] + active_count[spin_block] += ( + np.count_nonzero(geometry[:, spin_block], axis=0) + - geometry.shape[0] ) - value_count += output_np.shape[0] finally: self.mean, self.stddev = mean_backup, stddev_backup - if value_count == 0: + if not np.any(active_count > 0.0): return - feature_rms = np.sqrt(square_sum / float(value_count)) - if np.any(~np.isfinite(feature_rms)) or np.any(feature_rms <= self._STAT_EPS): + # A coordinate that never activates carries no information to + # precondition. That is the normal state of every spin coordinate on a + # demagnetized calibration corpus, which is exactly the corpus used to + # pretrain a model that is later fine-tuned on magnetic data, so it + # takes the identity preconditioner rather than an error. + measured = active_count > 0.0 + feature_rms = np.ones(geometry_dim, dtype=np.float64) + feature_rms[measured] = np.sqrt( + square_sum[measured] / active_count[measured], + ) + geometric = np.zeros(geometry_dim, dtype=bool) + geometric[: self.readout.get_dim_out()] = True + geometric[geometry_dim - 2 :] = True + degenerate = geometric & ( + ~np.isfinite(feature_rms) | (feature_rms <= self._STAT_EPS) + ) + if np.any(degenerate): raise ValueError( "DPA4C output calibration requires non-degenerate finite " - f"features, got RMS values {feature_rms.tolist()}" + f"geometric features, got RMS values {feature_rms.tolist()}" ) type_table = to_numpy_array(self.type_embedding.call())[: self.ntypes] target_rms = float(np.sqrt(np.mean(np.square(type_table, dtype=np.float64)))) @@ -1122,7 +1297,17 @@ def compute_input_stats( raise ValueError( f"DPA4C type embedding has a degenerate calibration RMS {target_rms}" ) - geometry_stddev = feature_rms / target_rms + # A coordinate earns a preconditioner only where its measured scale is + # meaningful. The geometric block is already required to be + # non-degenerate above; the spin block is not, because a corpus whose + # moments are uniformly weak drives the quartic spin coordinates to a + # vanishing root mean square, and dividing by it would hand them an + # unbounded gain. Those coordinates keep the identity instead, which is + # the same treatment a coordinate that never activates receives. + conditioned = measured & np.isfinite(feature_rms) + conditioned &= feature_rms > self._STAT_EPS + geometry_stddev = np.ones(geometry_dim, dtype=np.float64) + geometry_stddev[conditioned] = feature_rms[conditioned] / target_rms geometry_mean = np.zeros(geometry_dim, dtype=np.float64) # The two moment divisors are the only outputs carrying their @@ -1132,11 +1317,10 @@ def compute_input_stats( # standardized; every other coordinate keeps the shared RMS # preconditioner, whose zero mean the readout construction justifies. mass = slice(geometry_dim - 2, geometry_dim) - mass_mean = value_sum[mass] / float(value_count) + mass_count = active_count[mass] + mass_mean = value_sum[mass] / mass_count mass_stddev = np.sqrt( - np.maximum( - square_sum[mass] / float(value_count) - np.square(mass_mean), 0.0 - ) + np.maximum(square_sum[mass] / mass_count - np.square(mass_mean), 0.0) ) if np.any(mass_stddev <= self._STAT_EPS): raise ValueError( @@ -1154,6 +1338,105 @@ def compute_input_stats( PRECISION_DICT[self.precision] ) + def _calibration_frames(self, system: dict) -> list[dict]: + """Draw the calibration frames of one sampled system. + + The frames are taken on a linear index grid, so the draw is + deterministic and spreads over the whole system rather than over its + leading frames. + + Parameters + ---------- + system + Sampled system carrying ``coord``, ``atype``, an optional ``box``, + and, for a spin-conditioned descriptor, the per-atom moment under + either ``model_spin`` or ``spin``. + + Returns + ------- + list[dict] + One entry per drawn frame, each with a leading frame axis of + length one and a ``spin`` entry that is ``None`` for a spin-free + descriptor. + + Raises + ------ + ValueError + If a spin-conditioned descriptor is calibrated on a system that + carries no moment under either key. + """ + coord = to_numpy_array(system["coord"]) + nframes = coord.shape[0] + coord = np.reshape(coord, (nframes, -1, 3)) + atype = np.reshape(to_numpy_array(system["atype"]), (nframes, -1)) + box = system.get("box") + box = None if box is None else np.reshape(to_numpy_array(box), (nframes, -1)) + # The two keys are the two packings the training pipelines use: the + # native-spin route hands the moment through as ``spin``, while the + # virtual-atom route packs the model-facing arrays under a ``model_`` + # prefix so they survive next to the physical ones. + spin = ( + None + if self.spin is None + else self.require_spin(system.get("model_spin", system.get("spin"))) + ) + spin = ( + None if spin is None else np.reshape(to_numpy_array(spin), (nframes, -1, 3)) + ) + indices = np.linspace( + 0, + nframes - 1, + num=min(nframes, self._STAT_FRAMES_PER_SAMPLE), + dtype=np.int64, + ) + return [ + { + "coord": coord[index : index + 1], + "atype": atype[index : index + 1], + "box": None if box is None else box[index : index + 1], + "spin": None if spin is None else spin[index : index + 1], + } + for index in indices + ] + + def _measure_spin_reference(self, sampled: list[dict]) -> np.ndarray: + """Measure the per-type root-mean-square magnetic moment. + + The spin invariants are quadratic and quartic in the spin, so a + chemistry whose moments differ by a factor of three spreads the + quartic coordinates by two orders of magnitude. Rescaling the spin by + a per-type reference collapses that spread before the fixed diagonal + preconditioner sees it. The reference is a constant of the type, so + the rescaled spin stays linear in the input and every smoothness + property is preserved. + + Parameters + ---------- + sampled + Sampled training systems, each carrying a per-atom moment. + + Returns + ------- + numpy.ndarray + Strictly positive reference magnitudes with shape + ``(ntypes + 1,)``. A type that the sample never observes with a + finite moment keeps a unit reference, which leaves its spin + untouched; the estimator therefore never emits a zero. + """ + square_sum = np.zeros(self.ntypes + 1, dtype=np.float64) + count = np.zeros(self.ntypes + 1, dtype=np.float64) + for system in sampled: + for frame in self._calibration_frames(system): + atype = np.reshape(frame["atype"], (-1,)) + spin = np.reshape(frame["spin"], (-1, 3)) + magnitude = np.sum(np.square(spin, dtype=np.float64), axis=-1) + np.add.at(square_sum, atype, magnitude) + np.add.at(count, atype, 1.0) + reference = np.ones(self.ntypes + 1, dtype=np.float64) + observed = (count > 0.0) & (square_sum > count * self._STAT_EPS) + reference[observed] = np.sqrt(square_sum[observed] / count[observed]) + return reference + # === Serialization and neighbor statistics === def serialize(self) -> dict: @@ -1186,7 +1469,9 @@ def serialize(self) -> dict: "trainable": self.trainable, "type_map": self.type_map, "seed": self.seed, + "use_spin": self.use_spin, "spin": None, + "spin_channels": (None if self.spin is None else self.spin.serialize()), "type_embedding": self.type_embedding.serialize(), "radial_basis": self.radial_basis.serialize(), "radial_embedding": self.radial_embedding.serialize(), @@ -1237,6 +1522,7 @@ def deserialize(cls, data: dict) -> DescrptDPA4C: radial_mode_head = data.pop("radial_mode_head") pair_film = data.pop("pair_film") readout = data.pop("readout") + spin_channels = data.pop("spin_channels") obj = cls(**data) obj.type_embedding = SeZMTypeEmbedding.deserialize(type_embedding) @@ -1249,6 +1535,9 @@ def deserialize(cls, data: dict) -> DescrptDPA4C: ) obj.pair_film = OrderedPairFiLM.deserialize(pair_film) obj.readout = InvariantReadout.deserialize(readout) + obj.spin = ( + None if spin_channels is None else SpinChannels.deserialize(spin_channels) + ) obj.set_stat_mean_and_stddev( variables["mean"], variables["stddev"], @@ -1368,12 +1657,16 @@ def get_dim_out(self) -> int: ``K * (K + 1) // 2`` times the remaining rank; - three equal degrees contribute ``K * (K + 1) * (K + 2) // 6``. + A spin-conditioned descriptor appends the invariants of the spin + channels to the geometric block, ahead of the two divisors. + Returns ------- int Width of the invariant descriptor consumed by the fitting network. """ - return self.readout.get_dim_out() + self.channels + 2 + spin_dim = 0 if self.spin is None else self.spin.get_dim_out() + return self.readout.get_dim_out() + spin_dim + self.channels + 2 def get_dim_emb(self) -> int: """Return zero because fitting receives no equivariant channels.""" @@ -1384,9 +1677,18 @@ def mixed_types(self) -> bool: return True def has_message_passing(self) -> bool: - """Return whether source-node features are exchanged.""" + """Return whether source-node features are exchanged. + + The spin branch reads the raw source spin, which is a node input + rather than a derived feature, so a spin-conditioned descriptor + remains message-passing free. + """ return False + def supports_native_spin(self) -> bool: + """Return whether ``call_graph`` conditions on a per-node spin.""" + return self.spin is not None + def has_message_passing_across_ranks(self) -> bool: """Return whether intermediate halo communication is required.""" return False diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py index 4d5d2fbdbf..75b73b4a9c 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py @@ -21,18 +21,26 @@ from .readout import ( InvariantReadout, ) +from .spin import ( + NEIGHBOR_QUADRUPOLE_CHANNELS, + SpinChannels, + derive_spin_channels, +) __all__ = [ "MAX_ANGULAR_DEGREE", + "NEIGHBOR_QUADRUPOLE_CHANNELS", "BispectrumLayout", "InvariantReadout", "OrderedPairFiLM", + "SpinChannels", "build_angular_basis", "build_bispectrum_layout", "build_moment_indices", "degree_offsets", "derive_bispectrum_ranks", "derive_degree_channels", + "derive_spin_channels", "enumerate_degree_triples", "packed_l2_to_stf", ] diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py b/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py index af521eab84..f8be2e5095 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py @@ -5,16 +5,28 @@ annotations, ) +import math from typing import ( Any, ) import array_api_compat +import numpy as np from deepmd.dpmodel import ( DEFAULT_PRECISION, + PRECISION_DICT, NativeOP, ) +from deepmd.dpmodel.array_api import ( + xp_asarray_nodetach, +) +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -35,16 +47,37 @@ class OrderedPairFiLM(NativeOP): z_{ab}&=[T_a\Vert T_b],\\ h_{ab}&=\operatorname{SwiGLU}(z_{ab}W_{\rm in}),\\ - [s_{ab},d_{ab},u_{ab}]&=0.1\,h_{ab}W_{\rm out},\\ + [s_{ab},d_{ab},u_{ab},p_{ab},q_{ab}]&=0.1\,h_{ab}W_{\rm out},\\ \gamma_{ab}&=1+\tanh(s_{ab}),\\ \beta_{ab}&=T_a+T_b+\tanh(d_{ab}),\\ - U_{ab}&=\tanh(u_{ab}). + U_{ab}&=\tanh(u_{ab}),\\ + \gamma^{s}_{ab}&=\tanh(a^{\gamma}+p_{ab}),\qquad + \beta^{s}_{ab}=\tanh(a^{\beta}+q_{ab}). The network is evaluated over the finite type table rather than over graph - edges, so compressed inference stores only the three resulting tables. - Bounding every output keeps the cache well conditioned in ``float32``: + edges, so compressed inference stores only the resulting tables. Bounding + every output keeps the cache well conditioned in ``float32``: :math:`\gamma` stays in :math:`(0,2)`, and the residual parts of - :math:`\beta` and :math:`U` stay in :math:`(-1,1)`. + :math:`\beta`, :math:`U`, :math:`\gamma^{s}` and :math:`\beta^{s}` stay in + :math:`(-1,1)`. + + The spin scale :math:`\gamma^{s}` is signed, unlike its geometric + counterpart. The exchange interaction of an ordered pair may be either + ferromagnetic or antiferromagnetic, and the radial map is shared across + pairs, so the sign has to be available in the pair cache. + + That sign freedom rules out the structural anchor the geometric heads use + -- the constant one for :math:`\gamma`, the type embedding for + :math:`\beta` -- and the SwiGLU trunk is bias free, so without an anchor + of their own the two spin heads would emerge from a small centred product + and start several orders of magnitude below the geometric tables. The + descriptor calibration would then freeze a preconditioner at a scale the + first optimizer steps immediately leave. Both heads are therefore anchored + on a learned per-channel offset :math:`a`, initialized to + :math:`\pm\operatorname{artanh}` of :attr:`_SPIN_ANCHOR_MAGNITUDE` with an + independent random sign per channel. Every spin channel therefore starts + at that magnitude exactly, for any channel width and any seed, with its + sign left free. Parameters ---------- @@ -53,6 +86,9 @@ class OrderedPairFiLM(NativeOP): radial_modes Number :math:`R` of shared radial mode profiles each ordered pair mixes. Zero omits the mixing table. + spin_channels + Number :math:`C_s` of ordered spin scale and shift channels. Zero + omits the spin tables. precision Parameter precision. trainable @@ -63,15 +99,22 @@ class OrderedPairFiLM(NativeOP): Raises ------ ValueError - If ``channels`` is not positive or ``radial_modes`` is negative. + If ``channels`` is not positive, or if ``radial_modes`` or + ``spin_channels`` is negative. """ _OUTPUT_SCALE = 0.1 + #: Initial magnitude of both ordered spin tables. It is the same order as + #: the geometric scale and shift and stays well clear of the saturated + #: region of the bounding nonlinearity, whose derivative there is 0.75. + _SPIN_ANCHOR_MAGNITUDE = 0.5 + def __init__( self, channels: int, radial_modes: int = 0, + spin_channels: int = 0, *, precision: str = DEFAULT_PRECISION, trainable: bool = True, @@ -81,23 +124,41 @@ def __init__( raise ValueError(f"`channels` must be positive, got {channels}") if radial_modes < 0: raise ValueError(f"`radial_modes` must be non-negative, got {radial_modes}") + if spin_channels < 0: + raise ValueError( + f"`spin_channels` must be non-negative, got {spin_channels}" + ) self.channels = int(channels) self.radial_modes = int(radial_modes) + self.spin_channels = int(spin_channels) self.precision = str(precision) self.trainable = bool(trainable) input_dim = 2 * self.channels self.hidden_dim = resolve_swiglu_hidden_width(input_dim) - output_dim = self.channels * (2 + self.radial_modes) + output_dim = self.channels * (2 + self.radial_modes) + 2 * self.spin_channels self.network = SwiGLUMLP( [input_dim, self.hidden_dim, output_dim], output_scale=self._OUTPUT_SCALE, precision=self.precision, trainable=self.trainable, - seed=seed, + seed=child_seed(seed, 0), ) + if self.spin_channels == 0: + # The anchors exist only alongside the spin head they bias. + self.adam_spin_scale_anchor = None + self.adam_spin_shift_anchor = None + else: + rng = np.random.default_rng(child_seed(seed, 1)) + precision_dtype = PRECISION_DICT[self.precision.lower()] + offset = math.atanh(self._SPIN_ANCHOR_MAGNITUDE) + signs = rng.integers(0, 2, size=(2, self.spin_channels)) * 2.0 - 1.0 + self.adam_spin_scale_anchor = (offset * signs[0]).astype(precision_dtype) + self.adam_spin_shift_anchor = (offset * signs[1]).astype(precision_dtype) - def call(self, type_embedding: Any) -> tuple[Any, Any, Any | None]: - """Build the ordered scale, shift, and mixing tables. + def call( + self, type_embedding: Any + ) -> tuple[Any, Any, Any | None, Any | None, Any | None]: + """Build the ordered scale, shift, mixing, and spin tables. Parameters ---------- @@ -115,8 +176,15 @@ def call(self, type_embedding: Any) -> tuple[Any, Any, Any | None]: Ordered mode-mixing matrices with shape ``((T + 1) ** 2, channels, radial_modes)``, or ``None`` when ``radial_modes`` is zero. + spin_scale + Ordered spin scales with shape + ``((T + 1) ** 2, spin_channels)``, or ``None`` when + ``spin_channels`` is zero. + spin_shift + Ordered spin shifts with the same shape as ``spin_scale``. """ xp = array_api_compat.array_namespace(type_embedding) + device = array_api_compat.device(type_embedding) ntypes = type_embedding.shape[0] pair_shape = (ntypes, ntypes, self.channels) pair_input = xp.reshape( @@ -131,22 +199,46 @@ def call(self, type_embedding: Any) -> tuple[Any, Any, Any | None]: ) logits = self.network.call(pair_input) - # The output splits into the scale, the shift residual, and the - # flattened mixing matrix, in that order. + # The output splits into the scale, the shift residual, the flattened + # mixing matrix, and the two spin tables, in that order. shift_end = 2 * self.channels + mixing_end = shift_end + self.channels * self.radial_modes + spin_scale_end = mixing_end + self.spin_channels base_shift = xp.reshape( type_embedding[:, None, :] + type_embedding[None, :, :], (-1, self.channels), ) + + def anchored(anchor: Any, block: Any) -> Any: + """Bound one spin head around its learned per-channel offset.""" + return xp.tanh( + block + + xp_asarray_nodetach( + xp, + anchor, + dtype=block.dtype, + device=device, + )[None, :] + ) + return ( 1.0 + xp.tanh(logits[:, : self.channels]), base_shift + xp.tanh(logits[:, self.channels : shift_end]), None if self.radial_modes == 0 else xp.reshape( - xp.tanh(logits[:, shift_end:]), + xp.tanh(logits[:, shift_end:mixing_end]), (-1, self.channels, self.radial_modes), ), + None + if self.spin_channels == 0 + else anchored( + self.adam_spin_scale_anchor, + logits[:, mixing_end:spin_scale_end], + ), + None + if self.spin_channels == 0 + else anchored(self.adam_spin_shift_anchor, logits[:, spin_scale_end:]), ) def serialize(self) -> dict[str, Any]: @@ -155,16 +247,24 @@ def serialize(self) -> dict[str, Any]: Returns ------- dict[str, Any] - Versioned configuration and pair-encoder parameters. + Versioned configuration, pair-encoder parameters, and the two spin + anchors, which are present only for a spin-conditioned cache. """ return { "@class": "OrderedPairFiLM", "@version": 1, "channels": self.channels, "radial_modes": self.radial_modes, + "spin_channels": self.spin_channels, "precision": self.precision, "trainable": self.trainable, "network": self.network.serialize(), + "@variables": {} + if self.spin_channels == 0 + else { + "adam_spin_scale_anchor": to_numpy_array(self.adam_spin_scale_anchor), + "adam_spin_shift_anchor": to_numpy_array(self.adam_spin_shift_anchor), + }, } @classmethod @@ -191,6 +291,10 @@ def deserialize(cls, data: dict[str, Any]) -> OrderedPairFiLM: if data.pop("@class") != "OrderedPairFiLM": raise ValueError("Invalid serialized class for OrderedPairFiLM") network = data.pop("network") + variables = data.pop("@variables") obj = cls(**data) obj.network = SwiGLUMLP.deserialize(network) + precision_dtype = PRECISION_DICT[obj.precision.lower()] + for name, value in variables.items(): + setattr(obj, name, np.asarray(value, dtype=precision_dtype)) return obj diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py new file mode 100644 index 0000000000..ed936a1daf --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py @@ -0,0 +1,790 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Native per-atom spin channels for DPA4C. + +The magnetic moment :math:`\mathbf s` is an *axial* vector: under spatial +inversion the edge direction flips while the spin does not, and under time +reversal the spin flips while the geometry does not. Labelling every channel +by its angular degree :math:`\ell` and by its spin order :math:`\sigma` +(modulo two), and restricting every angular coupling to the genuine Gaunt +couplings the descriptor already uses (:math:`\ell_1+\ell_2+\ell_3` even), +gives + +.. math:: + + p=(-1)^{\ell}(-1)^{\sigma},\qquad t=(-1)^{\sigma}, + +so the two Z2 conditions decouple into the even-degree rule DPA4C already +enforces and one new rule: **only contractions of even total spin order may be +emitted**. The channel families below are laid out so that this rule is +realized by the block structure itself rather than by a per-entry filter. + +Five families are accumulated. + +============ ============ ============== ==================================== +Family :math:`\ell` :math:`\sigma` Content +============ ============ ============== ==================================== +``V`` 1 1 centre and neighbour spin vectors +``P`` 1 1 bond-projected neighbour spins +``Q`` 2 0 centre and neighbour spin quadrupoles +``M0`` 0 0 neighbour moment magnitude +``Mw`` 0 0 magnetic effective coordination +============ ============ ============== ==================================== + +``P`` is the one family that reads the edge direction. Each of its channels +accumulates :math:`(\hat{\mathbf s}_j\cdot\hat{\mathbf u}_{ij}) +\hat{\mathbf u}_{ij}`, a Cartesian vector carrying one moment and two factors +of the unit bond direction. Spatial inversion flips both direction factors and +leaves the moment alone, so the product is even; time reversal flips the +moment alone, so it is odd. That is exactly the grading of ``V``, which is why +the two occupy one block on the channel axis and the Gram of that block emits +every resulting invariant with no filtering rule of its own. Because ``P`` +depends on where a neighbour lies, the spin branch carries an angular +cotangent and its coordinate gradient flows through the direction as well as +through the distance. + +Their admissible contractions are the Gram of the joint ``V``/``P`` block +(Heisenberg exchange, symmetric anisotropic two-ion exchange, and the +collective anisotropies), the Gram of ``Q`` (biquadratic exchange), the cross +Gram of ``Q`` against the geometric degree-two moments (single-ion +anisotropy), and ``M0``/``Mw`` emitted directly. The cross Gram of the vector +block against the geometric degree-one moments has odd spin order and is +deliberately absent. + +Antisymmetric spin bilinears, and with them the Dzyaloshinskii-Moriya +interaction, remain unreachable at every order: the readout contracts channels +only through symmetric Grams and symmetric Gaunt couplings, so no +antisymmetric pairing of two moments is ever formed. + +``Mw`` is the one family that does not read the spin *value*: it is the +radially weighted count of neighbours that carry a magnetic degree of freedom +at all, and it is therefore nonzero even at vanishing spin. It is retained +because every other spin family vanishes with the moments, so without it the +readout cannot distinguish a neighbourhood with no magnetic species from a +magnetic neighbourhood that happens to be demagnetized -- the distinction that +separates a paramagnetic configuration from a non-magnetic one. Being +independent of the spin value, it contributes nothing to the magnetic force. +""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.array_api import ( + xp_asarray_nodetach, +) +from deepmd.dpmodel.common import ( + get_xp_precision, + to_numpy_array, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .geometry import ( + build_angular_basis, +) + +if TYPE_CHECKING: + from collections.abc import ( + Sequence, + ) + + from deepmd.dpmodel.array_api import ( + Array, + ) + +#: Number of quadrupole channels reduced over neighbours. One channel is +#: enough for the biquadratic exchange it enables, and each further channel +#: costs five moment accumulators against the three of a vector channel. +NEIGHBOR_QUADRUPOLE_CHANNELS = 1 + + +def derive_spin_channels(degree_channels: Sequence[int]) -> int: + """Derive the neighbour spin channel width. + + The spin channels index independent radial shapes of the effective + exchange function, and they read the leading channels of the same shared + radial map as the geometric degrees. The width is tied to the degree-two + width, which is the narrowest multi-channel degree of every supported + profile: that keeps the spin Gram, which grows quadratically, in + proportion to the rest of the profile, and it guarantees that every spin + channel addresses an already-evaluated radial channel. + + Parameters + ---------- + degree_channels + Channel widths for degrees zero through ``lmax``. The profile must + reach degree two, which every supported ``lmax`` does. + + Returns + ------- + int + Neighbour spin channel width :math:`C_s`. + + Raises + ------ + ValueError + If the degree profile does not reach degree two. + """ + if len(degree_channels) < 3: + raise ValueError( + "DPA4C native spin requires a degree profile reaching degree two, " + f"got {list(degree_channels)}" + ) + return int(degree_channels[2]) + + +class SpinChannels(NativeOP): + r"""Accumulate and contract the native spin channels of DPA4C. + + The per-atom spin enters through one conditioned node quantity + + .. math:: + + \hat{\mathbf s}_i=\frac{m_{a_i}}{s^{\mathrm{ref}}_{a_i}}\,\mathbf s_i, + + where :math:`m` is the per-type spin mask and :math:`s^{\mathrm{ref}}` the + per-type reference magnitude. The mask is multiplicative and is applied + once, so :math:`\partial E/\partial\mathbf s_i` vanishes identically -- at + every derivative order -- for a type that carries no magnetic degree of + freedom. Relying on the dataset convention :math:`\mathbf s=0` would not + achieve this, because the force loss differentiates the magnetic force + again and therefore probes the spin direction even where the value is + zero. + + Neighbour contributions share the radial map of the geometric branch and + add one ordered type-pair cache of their own: + + .. math:: + + \phi^{s}_{ij,c}=\chi_{ij}^2\bigl( + \gamma^{s}_{ab,c}g_c(\rho_{ij})+\beta^{s}_{ab,c}\bigr). + + The squared envelope matches the weight of every non-scalar geometric + moment, so the neighbour spin families share the existing normalizer + :math:`n^{(+)}` and need no neighbourhood mass of their own. Unlike the + geometric scale, :math:`\gamma^{s}` is signed: the exchange interaction of + an ordered pair may have either sign, and the shared radial map cannot + supply that sign per pair. + + Parameters + ---------- + ntypes + Number of real atom types. + degree_channels + Channel widths for degrees zero through ``lmax``. + use_spin + Per-type flags marking which atom types carry a magnetic moment. Every + weight is sized by the type count rather than by the magnetic subset, + and the flags only build the per-type gate, which is derived from the + configuration rather than serialized. An all-false mask is therefore a + supported configuration: it contributes an identical zero and keeps the + full set of spin parameters, which is what a spin-free pretrain needs in + order to declare its magnetic types at fine-tune time. + precision + Parameter precision. + trainable + Whether the on-site weights receive optimizer updates. + seed + Random seed. + + Raises + ------ + ValueError + If ``use_spin`` does not have one entry per real atom type. + """ + + def __init__( + self, + ntypes: int, + degree_channels: Sequence[int], + use_spin: Sequence[bool], + *, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + if len(use_spin) != int(ntypes): + raise ValueError( + f"`use_spin` must contain {int(ntypes)} entries, got {len(use_spin)}" + ) + self.ntypes = int(ntypes) + self.degree_channels = [int(width) for width in degree_channels] + self.use_spin = [bool(flag) for flag in use_spin] + self.precision = str(precision) + self.trainable = bool(trainable) + self.spin_channels = derive_spin_channels(self.degree_channels) + precision_dtype = PRECISION_DICT[self.precision.lower()] + + # === Per-type spin gate === + # Deterministic from the configuration, so it is rebuilt rather than + # serialized. The trailing row is the padding type. + self.spin_mask = np.asarray( + [1.0 if flag else 0.0 for flag in self.use_spin] + [0.0], + dtype=precision_dtype, + ) + # Per-type reference magnitude in the units of the dataset. Measured + # by the descriptor calibration and therefore persistent; a unit + # reference leaves the raw spin untouched. + self.spin_reference = np.ones(self.ntypes + 1, dtype=precision_dtype) + + # === On-site weights === + # One channel each, so a per-type scalar. The invariants they enter + # are linear in this weight, and the fitting network already receives + # the centre type embedding, so additional on-site channels would only + # rescale the same quantities. + rng = np.random.default_rng(child_seed(seed, 0)) + self.adam_spin_vector_weight = rng.normal( + 0.0, 1.0, size=(self.ntypes + 1,) + ).astype(precision_dtype) + self.adam_spin_quadrupole_weight = rng.normal( + 0.0, 1.0, size=(self.ntypes + 1,) + ).astype(precision_dtype) + + # Isometric half-vectorization of the two spin Grams. The quadrupole + # block drops entry zero, its on-site self-term: the harmonic blocks + # are homogeneous, so |B_2(s)|^2 = |s|^4 exactly and that entry is a + # per-type constant times the square of the on-site self-term the + # vector block already emits. + self.vector_gram_index, self.vector_gram_scale = _half_gram_layout( + self.vector_width, + precision_dtype, + ) + quadrupole_index, quadrupole_scale = _half_gram_layout( + self.quadrupole_width, + precision_dtype, + ) + self.quadrupole_gram_index = quadrupole_index[1:] + self.quadrupole_gram_scale = quadrupole_scale[1:] + + def call(self, spin_moments: Array, degree_two: Array) -> Array: + r"""Contract the spin moments into invariants of even spin order. + + Parameters + ---------- + spin_moments + Flat spin moments with shape ``(N, moment_width)``, normalized and + concatenated by the descriptor. + degree_two + Geometric degree-two moments with shape ``(N, 5, C_2)``. + + Returns + ------- + Array + Invariant spin features with shape ``(N, get_dim_out())``. + """ + xp = array_api_compat.array_namespace(spin_moments) + magnitude, coordination, vector, quadrupole = self.split(spin_moments, xp) + return xp.concat( + [ + _half_gram( + vector, + self.vector_gram_index, + self.vector_gram_scale, + xp, + ), + _half_gram( + quadrupole, + self.quadrupole_gram_index, + self.quadrupole_gram_scale, + xp, + ), + # Cross Gram against the geometric degree-two moments. Both + # factors have even spin order, so the product is admissible; + # with a unit direction it evaluates to the single-ion + # anisotropy sum over neighbours. + xp.reshape( + xp.matmul(xp.permute_dims(quadrupole, (0, 2, 1)), degree_two), + (spin_moments.shape[0], -1), + ), + magnitude, + coordination, + ], + axis=-1, + ) + + def split(self, spin_moments: Array, xp: Any) -> tuple[Array, Array, Array, Array]: + """Split the flat spin moments into their five families. + + Parameters + ---------- + spin_moments + Flat spin moments with shape ``(N, moment_width)``. + xp + Array namespace associated with ``spin_moments``. + + Returns + ------- + magnitude + Neighbour moment magnitudes with shape ``(N, spin_channels)``. + coordination + Magnetic effective coordination with shape ``(N, spin_channels)``. + vector + Joint degree-one spin block with shape ``(N, 3, vector_width)``, + holding the on-site moment, the ``V`` channels and the ``P`` + channels in that order. + quadrupole + Degree-two spin block with shape ``(N, 5, quadrupole_width)``. + """ + nodes = spin_moments.shape[0] + channels = self.spin_channels + neighbor_quadrupole = NEIGHBOR_QUADRUPOLE_CHANNELS + offset = 0 + magnitude = spin_moments[:, offset : offset + channels] + offset += channels + coordination = spin_moments[:, offset : offset + channels] + offset += channels + neighbor_vector = xp.reshape( + spin_moments[:, offset : offset + 3 * channels], + (nodes, 3, channels), + ) + offset += 3 * channels + neighbor_bond = xp.reshape( + spin_moments[:, offset : offset + 3 * channels], + (nodes, 3, channels), + ) + offset += 3 * channels + neighbor_tensor = xp.reshape( + spin_moments[:, offset : offset + 5 * neighbor_quadrupole], + (nodes, 5, neighbor_quadrupole), + ) + offset += 5 * neighbor_quadrupole + onsite_vector = spin_moments[:, offset : offset + 3] + offset += 3 + onsite_tensor = spin_moments[:, offset : offset + 5] + # The on-site channel leads each block so that its Gram entries + # against the neighbour channels, which carry the two-body physics, + # occupy the first row of the upper triangle. + return ( + magnitude, + coordination, + xp.concat( + [onsite_vector[:, :, None], neighbor_vector, neighbor_bond], + axis=-1, + ), + xp.concat([onsite_tensor[:, :, None], neighbor_tensor], axis=-1), + ) + + def conditioned_spin(self, spin: Array, atype: Array) -> Array: + r"""Apply the per-type spin mask and reference magnitude. + + Parameters + ---------- + spin + Per-node spin vectors with shape ``(N, 3)``. + atype + Flat node types with shape ``(N,)``. Padding nodes use type index + ``ntypes`` and select the zero row. + + Returns + ------- + Array + Conditioned spin :math:`\hat{\mathbf s}` with shape + ``(N, 3)``, exactly zero for non-magnetic and padding types. + """ + xp = array_api_compat.array_namespace(spin) + device = array_api_compat.device(spin) + dtype = get_xp_precision(xp, self.precision) + # The gate and the reference collapse into one multiplicative table, + # which costs one gather instead of two. Every reference entry is + # strictly positive -- the estimator seeds the table with ones and + # overwrites only strictly positive measurements, and both setters + # enforce it -- so the quotient is finite for every type, including + # the non-magnetic and padding rows whose gate is zero. + gate = xp_asarray_nodetach(xp, self.spin_mask[...], device=device) + reference = xp_asarray_nodetach(xp, self.spin_reference[...], device=device) + weight = xp.astype(gate / reference, dtype) + index = xp.astype(atype, xp.int64) + return xp.astype(spin, dtype) * xp.take(weight, index, axis=0)[:, None] + + def onsite_payload(self, conditioned_spin: Array, atype: Array) -> Array: + r"""Build the node-local on-site spin moments. + + The centre spin enters as the degree-one vector + :math:`\lambda_{a}\hat{\mathbf s}_i` and the degree-two quadrupole + :math:`\mu_{a}B_2(\hat{\mathbf s}_i)`. Both are polynomial in the + spin, hence smooth at :math:`\hat{\mathbf s}=0`, and both are node + local: they are written after the destination reduction and carry no + neighbourhood normalizer, so the invariants they enter contain exactly + one factor of :math:`n^{(+)}` from their neighbour partner. + + Parameters + ---------- + conditioned_spin + Conditioned spin with shape ``(N, 3)``. + atype + Flat node types with shape ``(N,)``. + + Returns + ------- + Array + On-site payload with shape ``(N, node_width)``. + """ + xp = array_api_compat.array_namespace(conditioned_spin) + device = array_api_compat.device(conditioned_spin) + index = xp.astype(atype, xp.int64) + vector_weight = xp.take( + xp_asarray_nodetach(xp, self.adam_spin_vector_weight, device=device), + index, + axis=0, + ) + quadrupole_weight = xp.take( + xp_asarray_nodetach(xp, self.adam_spin_quadrupole_weight, device=device), + index, + axis=0, + ) + quadrupole = build_angular_basis(conditioned_spin, 2)[:, 4:9] + return xp.concat( + [ + conditioned_spin * vector_weight[:, None], + quadrupole * quadrupole_weight[:, None], + ], + axis=-1, + ) + + def edge_payload( + self, + conditioned_spin: Array, + atype: Array, + source: Array, + direction: Array, + radial: Array, + envelope: Array, + pair_scale: Array, + pair_shift: Array, + pair_index: Array, + ) -> Array: + r"""Build the per-edge spin payload reduced over neighbours. + + Every family shares one edge amplitude, the spin counterpart of the + geometric :math:`\phi_{ij}` + + .. math:: + + \phi^{s}_{ij,c}=\chi_{ij}^2\bigl( + \gamma^{s}_{ab,c}g_c(\rho_{ij})+\beta^{s}_{ab,c}\bigr), + + so the radial table is read once for the whole spin branch. The + bond-projected family additionally contracts the unit edge direction, + + .. math:: + + \mathbf P_{ij,c}=\phi^{s}_{ij,c} + (\hat{\mathbf s}_j\cdot\hat{\mathbf u}_{ij})\hat{\mathbf u}_{ij}, + + which is what makes the symmetric anisotropic exchange representable. + + Parameters + ---------- + conditioned_spin + Conditioned spin with shape ``(N, 3)``. + atype + Flat node types with shape ``(N,)``, read for the neighbour spin + gate of the magnetic-coordination family. + source + Source node index of each edge with shape ``(E,)``. + direction + Regularized unit edge directions with shape ``(E, 3)``. + radial + Shared radial map with shape ``(E, channels)``; the leading + ``spin_channels`` columns are read. + envelope + Masked C3 envelope with shape ``(E,)``. + pair_scale + Ordered spin scales with shape ``((ntypes + 1) ** 2, spin_channels)``. + pair_shift + Ordered spin shifts with the same shape. + pair_index + Ordered type-pair index of each edge with shape ``(E,)``. + + Returns + ------- + Array + Edge payload with shape ``(E, edge_width)``. + """ + xp = array_api_compat.array_namespace(conditioned_spin) + device = array_api_compat.device(conditioned_spin) + channels = self.spin_channels + neighbor_spin = xp.take(conditioned_spin, source, axis=0) # (E, 3) + scale = xp.take(pair_scale, pair_index, axis=0) # (E, Cs) + shift = xp.take(pair_shift, pair_index, axis=0) # (E, Cs) + spin_amplitude = (radial[:, :channels] * scale + shift) * (envelope * envelope)[ + :, None + ] + neighbor_gate = xp.take( + xp.take( + xp_asarray_nodetach(xp, self.spin_mask, device=device), + xp.astype(atype, xp.int64), + axis=0, + ), + source, + axis=0, + )[:, None] # (E, 1) + + magnitude = xp.sum(neighbor_spin * neighbor_spin, axis=-1, keepdims=True) + # Component of the neighbour moment along the bond, carried back as a + # vector so that the block Gram turns it into the bond-resolved + # invariants. The masked envelope already zeroes excluded edges, so the + # unmasked direction cannot leak through the amplitude. + bond_spin = direction * xp.sum( + neighbor_spin * direction, + axis=-1, + keepdims=True, + ) # (E, 3) + quadrupole = build_angular_basis(neighbor_spin, 2)[:, 4:9] # (E, 5) + return xp.concat( + [ + spin_amplitude * magnitude, + spin_amplitude * neighbor_gate, + xp.reshape( + neighbor_spin[:, :, None] * spin_amplitude[:, None, :], + (-1, 3 * channels), + ), + xp.reshape( + bond_spin[:, :, None] * spin_amplitude[:, None, :], + (-1, 3 * channels), + ), + xp.reshape( + quadrupole[:, :, None] + * spin_amplitude[:, None, :NEIGHBOR_QUADRUPOLE_CHANNELS], + (-1, 5 * NEIGHBOR_QUADRUPOLE_CHANNELS), + ), + ], + axis=-1, + ) + + @property + def vector_width(self) -> int: + r"""Return the channel width of the joint degree-one spin block. + + The block holds the on-site moment, the :math:`C_s` isotropic + neighbour channels ``V`` and the :math:`C_s` bond-projected neighbour + channels ``P``. ``P`` is given the full spin width rather than a + narrower one because the readout is linear in the emitted Gram + entries: the effective radial profile of an interaction is the span of + the channel amplitudes the fitting network mixes, so a narrower ``P`` + would confine the symmetric anisotropic exchange to a smaller function + space than the isotropic exchange beside it and than the single-ion + anisotropy, which already reaches :math:`C_2=C_s` geometric channels + through its cross Gram. Reusing the leading :math:`C_s` amplitudes + costs no extra radial evaluation. + """ + return 1 + 2 * self.spin_channels + + @property + def quadrupole_width(self) -> int: + """Return the channel width of the degree-two spin block.""" + return 1 + NEIGHBOR_QUADRUPOLE_CHANNELS + + @property + def edge_width(self) -> int: + """Return the width of the per-edge payload reduced over neighbours. + + The layout is ``[M0, Mw, V_neighbor, P_neighbor, Q_neighbor]`` with + the harmonic component as the outer axis of each non-scalar family, + matching the geometric moment convention. + """ + return ( + 2 * self.spin_channels + + 6 * self.spin_channels + + 5 * NEIGHBOR_QUADRUPOLE_CHANNELS + ) + + @property + def node_width(self) -> int: + """Return the width of the node-local on-site payload ``[V_o, Q_o]``.""" + return 3 + 5 + + @property + def moment_width(self) -> int: + """Return the total spin moment width appended to the flat state.""" + return self.edge_width + self.node_width + + def get_dim_out(self) -> int: + """Return the invariant width contributed by the spin channels. + + Returns + ------- + int + Upper triangles of the two spin Grams, the full cross Gram against + the geometric degree-two moments, and the two scalar families. + """ + return ( + int(self.vector_gram_index.shape[0]) + + int(self.quadrupole_gram_index.shape[0]) + + self.quadrupole_width * self.degree_channels[2] + + 2 * self.spin_channels + ) + + def serialize(self) -> dict[str, Any]: + """Serialize the spin channels. + + Returns + ------- + dict[str, Any] + Versioned configuration and persistent arrays. The per-type mask + is deterministic from ``use_spin`` and is therefore rebuilt. + """ + return { + "@class": "SpinChannels", + "@version": 1, + "ntypes": self.ntypes, + "degree_channels": list(self.degree_channels), + "use_spin": list(self.use_spin), + "precision": self.precision, + "trainable": self.trainable, + "@variables": { + "spin_reference": to_numpy_array(self.spin_reference), + "adam_spin_vector_weight": to_numpy_array(self.adam_spin_vector_weight), + "adam_spin_quadrupole_weight": to_numpy_array( + self.adam_spin_quadrupole_weight + ), + }, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> SpinChannels: + """Deserialize a :class:`SpinChannels`. + + Parameters + ---------- + data + Versioned dictionary produced by :meth:`serialize`. + + Returns + ------- + SpinChannels + Reconstructed spin channels. + + Raises + ------ + ValueError + If the payload does not describe a :class:`SpinChannels`. + """ + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + if data.pop("@class") != "SpinChannels": + raise ValueError("Invalid serialized class for SpinChannels") + variables = data.pop("@variables") + obj = cls(**data) + obj.set_variables(variables) + return obj + + def set_variables(self, variables: dict[str, Any]) -> None: + """Restore the persistent arrays. + + Parameters + ---------- + variables + Mapping produced by the ``@variables`` block of :meth:`serialize`. + """ + precision_dtype = PRECISION_DICT[self.precision.lower()] + # Routed through the setter so a restored checkpoint cannot weaken the + # strictly-positive reference invariant that ``conditioned_spin`` relies on. + self.set_spin_reference(variables["spin_reference"]) + self.adam_spin_vector_weight = np.asarray( + variables["adam_spin_vector_weight"], dtype=precision_dtype + ) + self.adam_spin_quadrupole_weight = np.asarray( + variables["adam_spin_quadrupole_weight"], dtype=precision_dtype + ) + + def set_spin_reference(self, reference: np.ndarray) -> None: + """Store the per-type reference magnitudes. + + Parameters + ---------- + reference + Reference magnitudes with shape ``(ntypes + 1,)`` in the units of + the dataset spin. Every entry must be strictly positive, including + the non-magnetic and padding rows, whose reference is unused but + still divides the zero gate. + + Raises + ------ + ValueError + If the shape is wrong or any entry is not strictly positive. + """ + reference = np.asarray(to_numpy_array(reference), dtype=np.float64) + if reference.shape != (self.ntypes + 1,): + raise ValueError( + "DPA4C spin reference must have shape " + f"{(self.ntypes + 1,)}, got {reference.shape}" + ) + if not np.all(np.isfinite(reference)) or np.any(reference <= 0.0): + raise ValueError( + "DPA4C spin reference must be finite and strictly positive, " + f"got {reference.tolist()}" + ) + self.spin_reference = reference.astype(PRECISION_DICT[self.precision.lower()]) + + +def _half_gram_layout( + width: int, + precision_dtype: np.dtype, +) -> tuple[np.ndarray, np.ndarray]: + """Build the isometric upper-triangular Gram index and scale. + + The entries are ordered row by row, so entry zero is always the self-term + of the leading channel and a caller that does not emit it drops the first + element of both arrays. + + Parameters + ---------- + width + Channel width of the block. + precision_dtype + Element type of the emitted scale. + + Returns + ------- + index + Flattened Gram coordinates of the upper triangle. + scale + Matching isometric scales, ``sqrt(2)`` off the diagonal. + """ + row, column = np.triu_indices(int(width)) + return ( + (row * int(width) + column).astype(np.int64), + np.where(row == column, 1.0, math.sqrt(2.0)).astype(precision_dtype), + ) + + +def _half_gram( + block: Array, + index: np.ndarray, + scale: np.ndarray, + xp: Any, +) -> Array: + """Return the Frobenius-isometric upper-triangular channel Gram.""" + device = array_api_compat.device(block) + width = block.shape[-1] + gram = xp.reshape( + xp.matmul(xp.permute_dims(block, (0, 2, 1)), block), + (block.shape[0], width * width), + ) + return ( + xp.take( + gram, + xp_asarray_nodetach(xp, index, device=device), + axis=1, + ) + * xp_asarray_nodetach(xp, scale, device=device)[None, :] + ) diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index b0a77e22fc..f587cc2906 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -344,6 +344,21 @@ def enable_compression( """ raise NotImplementedError("This descriptor doesn't support compression!") + def compression_needs_min_nbor_dist(self) -> bool: + """Whether :meth:`enable_compression` consumes ``min_nbor_dist``. + + Returns + ------- + bool + Concrete default ``True``: a tabulated embedding starts its + table at the shortest distance the training data contains, so + the caller must measure it first. ``False`` for descriptors + whose table domain is fixed analytically; the caller may then + skip the neighbor-statistics pass, which is a dense all-pairs + computation over the training data. + """ + return True + @abstractmethod def fwd( self, diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b6f765c22a..bee5b79215 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -271,6 +271,10 @@ def enable_compression( check_frequency, ) + def compression_needs_min_nbor_dist(self) -> bool: + """Delegates to the atomic model.""" + return bool(self.atomic_model.compression_needs_min_nbor_dist()) + def call_common( self, coord: Array, diff --git a/deepmd/infer/model_test/ener.py b/deepmd/infer/model_test/ener.py index 6c02b1566e..985fbdcf72 100644 --- a/deepmd/infer/model_test/ener.py +++ b/deepmd/infer/model_test/ener.py @@ -268,7 +268,7 @@ def _write_energy_test_details( save_txt_file( detail_path.with_suffix(".s.out"), ps, - header=f"{system} (eV/Å^3): data_sxx data_sxy data_sxz data_syx " + header=f"{system} (eV/ų): data_sxx data_sxy data_sxz data_syx " "data_syy data_syz data_szx data_szy data_szz pred_sxx pred_sxy pred_sxz " "pred_syx pred_syy pred_syz pred_szx pred_szy pred_szz", append=append_detail, @@ -360,8 +360,8 @@ class EnerTester(ModelTester): ("rmse_v", "Virial RMSE : {} eV"), ("mae_va", "Virial MAE/Natoms : {} eV"), ("rmse_va", "Virial RMSE/Natoms : {} eV"), - ("mae_s", "Stress MAE : {} eV/Å^3"), - ("rmse_s", "Stress RMSE : {} eV/Å^3"), + ("mae_s", "Stress MAE : {} eV/ų"), + ("rmse_s", "Stress RMSE : {} eV/ų"), ("mae_ae", "Atomic ener MAE : {} eV"), ("rmse_ae", "Atomic ener RMSE : {} eV"), ("mae_h", "Hessian MAE : {} eV/Å^2"), @@ -513,7 +513,7 @@ def evaluate_chunk( if reports_virial: if shared_metrics.virial is None or shared_metrics.virial_per_atom is None: raise RuntimeError("Virial metrics are unavailable for dp test.") - # Stress sigma = -virial / volume, in eV/Å^3 (tensile-positive + # Stress sigma = -virial / volume, in eV/ų (tensile-positive # convention). volume = np.abs(np.linalg.det(box.reshape([nframes, 3, 3]))).reshape( [nframes, 1] diff --git a/deepmd/kernels/cuda/dpa1/canonical.py b/deepmd/kernels/cuda/dpa1/canonical.py index 113066e2a1..54c0c8b2ea 100644 --- a/deepmd/kernels/cuda/dpa1/canonical.py +++ b/deepmd/kernels/cuda/dpa1/canonical.py @@ -436,13 +436,14 @@ def dpa1_canonical_compress_energy_force( float(se.env_protection), float(se.nnei), ) - force, atom_virial, virial = canonical_edge_force_virial( + force, atom_virial, virial, _ = canonical_edge_force_virial( edge_gradient, graph.edge_vec, graph.destination_row_ptr, graph.source_row_ptr, graph.source_order, graph.n_node, + graph.edge_vec.new_zeros(0, 3), atype.shape[0], do_atomic_virial, ) diff --git a/deepmd/kernels/cuda/dpa1/graph_compress.py b/deepmd/kernels/cuda/dpa1/graph_compress.py index a05f83b56a..8a62384ebe 100644 --- a/deepmd/kernels/cuda/dpa1/graph_compress.py +++ b/deepmd/kernels/cuda/dpa1/graph_compress.py @@ -1012,7 +1012,7 @@ def dpa1_graph_compress_energy_force( float(se.env_protection), float(se.nnei), ) - force, atom_virial, virial = edge_force_virial( + force, atom_virial, virial, _ = edge_force_virial( edge_gradient, edge_vec, graph.edge_index, @@ -1022,6 +1022,7 @@ def dpa1_graph_compress_energy_force( graph.source_order, graph.source_row_ptr, graph.n_node, + edge_vec.new_zeros(0, 3), node_capacity, do_atomic_virial, ) diff --git a/deepmd/kernels/cuda/dpa1/graph_energy_force.py b/deepmd/kernels/cuda/dpa1/graph_energy_force.py index 168c7f3ca8..d3e627053c 100644 --- a/deepmd/kernels/cuda/dpa1/graph_energy_force.py +++ b/deepmd/kernels/cuda/dpa1/graph_energy_force.py @@ -266,7 +266,7 @@ def _cpu( # ``_fake`` and the CUDA operator; the sub-operator would otherwise return # fp64 whenever the edge inputs are fp64. fprec = w1.dtype - force, atom_virial, virial = torch.ops.deepmd.edge_force_virial( + force, atom_virial, virial, _ = torch.ops.deepmd.edge_force_virial( g_e.to(fprec), edge_vec.to(fprec), edge_index, @@ -276,6 +276,7 @@ def _cpu( source_order, source_row_ptr, n_node, + edge_vec.to(fprec).new_zeros(0, 3), node_capacity, do_atomic_virial, ) diff --git a/deepmd/kernels/cuda/dpa4c/canonical.py b/deepmd/kernels/cuda/dpa4c/canonical.py index 98cc400724..25504227d5 100644 --- a/deepmd/kernels/cuda/dpa4c/canonical.py +++ b/deepmd/kernels/cuda/dpa4c/canonical.py @@ -85,6 +85,9 @@ def _forward_fake( coupling_value: torch.Tensor, output_mean: torch.Tensor, output_inv_std: torch.Tensor, + spin: torch.Tensor, + spin_pair: torch.Tensor, + spin_type: torch.Tensor, lmax: int, table_stride: float, table_max: float, @@ -104,6 +107,8 @@ def _forward_fake( coupling_value, output_mean, output_inv_std, + spin_pair, + spin_type, table_stride, table_max, rcut, @@ -114,7 +119,9 @@ def _forward_fake( descriptor_profile, ) - profile = descriptor_profile(int(type_embedding.shape[1]), int(lmax)) + profile = descriptor_profile( + int(type_embedding.shape[1]), int(lmax), spin.numel() != 0 + ) nodes = atype.shape[0] descriptor = edge_vec.new_empty(nodes, profile.output_width, dtype=torch.float32) state = edge_vec.new_empty(nodes, profile.state_width, dtype=torch.float32) @@ -126,9 +133,9 @@ def _backward_fake( state: torch.Tensor, edge_vec: torch.Tensor, *args: Any, -) -> torch.Tensor: - del descriptor_gradient, state, args - return torch.empty_like(edge_vec) +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del descriptor_gradient, state + return _backward_shapes(edge_vec, args[13]) def _backward_inplace_fake( @@ -136,9 +143,31 @@ def _backward_inplace_fake( state: torch.Tensor, edge_vec: torch.Tensor, *args: Any, -) -> torch.Tensor: - del descriptor_gradient, state, args - return torch.empty_like(edge_vec) +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del descriptor_gradient, state + return _backward_shapes(edge_vec, args[13]) + + +def _backward_shapes( + edge_vec: torch.Tensor, + spin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return the three backward outputs of one compact call, unpopulated. + + Each absent output is allocated separately, because the schema declares + three unannotated results and two of them may not share storage. + """ + if spin.numel() == 0: + return ( + torch.empty_like(edge_vec), + edge_vec.new_empty((0,), dtype=torch.float32), + edge_vec.new_empty((0,), dtype=torch.float32), + ) + return ( + torch.empty_like(edge_vec), + spin.new_empty((spin.shape[0], 3)), + torch.empty_like(edge_vec), + ) def _energy_gradient_fake( @@ -147,11 +176,21 @@ def _energy_gradient_fake( destination_row_ptr: torch.Tensor, atype: torch.Tensor, *args: Any, -) -> tuple[torch.Tensor, torch.Tensor]: - del source, destination_row_ptr, args +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + del source, destination_row_ptr + spin = args[10] + has_spin = spin.numel() != 0 + # Each absent output is allocated separately, because the schema declares + # four unannotated results and no two of them may share storage. return ( edge_vec.new_empty(atype.shape[0], 1, dtype=torch.float64), torch.empty_like(edge_vec), + edge_vec.new_empty(atype.shape[0], 3, dtype=torch.float32) + if has_spin + else edge_vec.new_empty((0,), dtype=torch.float32), + torch.empty_like(edge_vec) + if has_spin + else edge_vec.new_empty((0,), dtype=torch.float32), ) @@ -172,8 +211,10 @@ def _cpu_energy_gradient(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: gradient = fitting_backward( seed.reshape(-1, 1), saved, ws, bs, resnets, w_head, act ) - edge_gradient = _cpu_backward(gradient, state, *descriptor_args) - return energy, edge_gradient + edge_gradient, spin_gradient, edge_spin_gradient = _cpu_backward( + gradient, state, *descriptor_args + ) + return energy, edge_gradient, spin_gradient, edge_spin_gradient def _generic_topology( @@ -211,7 +252,7 @@ def _generic_topology( #: Leading arguments of the fused operator that describe the descriptor: #: ``edge_vec`` plus the compact topology, the compression artifacts and the #: six trailing geometry scalars. -_DESCRIPTOR_ARGUMENT_COUNT = 20 +_DESCRIPTOR_ARGUMENT_COUNT = 23 def _cpu_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: @@ -310,7 +351,15 @@ def dpa4c_canonical_compress_energy_force( ownership: torch.Tensor, atom_bias: torch.Tensor, do_atomic_virial: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + spin: torch.Tensor | None = None, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: """Evaluate compressed DPA4C from a compact canonical edge stream. The compact ABI carries only source indices and CSR row pointers, so the @@ -335,6 +384,9 @@ def dpa4c_canonical_compress_energy_force( Combined atomic energy bias with shape ``(ntypes,)`` in eV. do_atomic_virial Whether to return per-node virials. + spin + Per-node magnetic moments with shape ``(N, 3)`` for a spin-conditioned + descriptor, or ``None``. Returns ------- @@ -348,6 +400,11 @@ def dpa4c_canonical_compress_energy_force( Per-frame virial with shape ``(F, 3, 3)`` in eV, fp32. atom_virial Per-node virial with shape ``(N, 3, 3)`` in eV, or an empty tensor. + force_mag + Per-node magnetic force with shape ``(N, 3)``, or an empty tensor. The + on-site part closes inside the fused operator; the neighbour part is + emitted per edge and reduced onto source nodes here, where the source + CSR is in scope. Raises ------ @@ -383,36 +440,43 @@ def dpa4c_canonical_compress_energy_force( raise ValueError("model is not eligible for compact canonical DPA4C inference") network = fitting_operator_arguments(fitting) - atom_energy_raw, edge_gradient = ( - torch.ops.deepmd.dpa4c_canonical_compress_energy_gradient( - graph.edge_vec, - graph.source, - graph.destination_row_ptr, - atype, - *compressed_operator_arguments(descriptor), - int(descriptor.lmax), - *descriptor._compression_scalars, - network.weights, - network.biases, - network.residuals, - network.head_weight, - network.head_bias, - atom_bias.to(torch.float64).contiguous(), - network.activation, - ownership.to(torch.float64).reshape(-1).contiguous(), - node_tile(), - ) + ( + atom_energy_raw, + edge_gradient, + spin_gradient, + edge_spin_gradient, + ) = torch.ops.deepmd.dpa4c_canonical_compress_energy_gradient( + graph.edge_vec, + graph.source, + graph.destination_row_ptr, + atype, + *compressed_operator_arguments(descriptor, spin), + int(descriptor.lmax), + *descriptor._compression_scalars, + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.head_bias, + atom_bias.to(torch.float64).contiguous(), + network.activation, + ownership.to(torch.float64).reshape(-1).contiguous(), + node_tile(), ) atom_energy = atom_energy_raw * ownership[:, None].to(atom_energy_raw.dtype) energy = frame_scalar_sum(atom_energy, graph.n_node) - force, atom_virial, virial = canonical_edge_force_virial( + # The magnetic cotangent is reduced onto its source nodes by the force + # assembly, which already walks that grouping. + force, atom_virial, virial, source_spin = canonical_edge_force_virial( edge_gradient, graph.edge_vec, graph.destination_row_ptr, graph.source_row_ptr, graph.source_order, graph.n_node, + edge_spin_gradient, atype.shape[0], do_atomic_virial, ) - return energy, atom_energy, force, virial, atom_virial + force_mag = spin_gradient if spin is None else -(spin_gradient + source_spin) + return energy, atom_energy, force, virial, atom_virial, force_mag diff --git a/deepmd/kernels/cuda/dpa4c/graph_compress.py b/deepmd/kernels/cuda/dpa4c/graph_compress.py index f35f25f13f..bcb3d845b6 100644 --- a/deepmd/kernels/cuda/dpa4c/graph_compress.py +++ b/deepmd/kernels/cuda/dpa4c/graph_compress.py @@ -48,6 +48,7 @@ build_bispectrum_layout, derive_bispectrum_ranks, derive_degree_channels, + derive_spin_channels, packed_l2_to_stf, ) @@ -63,6 +64,7 @@ "fitting_energy_and_gradient", "mega_eligible", "op_available", + "reduce_edge_spin_gradient", ] SUPPORTED_CHANNELS = (8, 16, 32, 64, 128) @@ -107,6 +109,12 @@ def ef_op_available() -> bool: def mega_eligible(descriptor: Any) -> bool: """Return whether the descriptor has a compiled fp32 specialization. + Native spin is a compiled variant of the same kernels rather than a + separate operator, and every spin width follows from the channel count, so + a spin-conditioned descriptor is eligible on exactly the conditions a + spin-free one is. Each condition below is a width the compiled operator + specializes on. + Parameters ---------- descriptor @@ -160,15 +168,29 @@ class DescriptorProfile: output_width: int gram_base: int bispectrum_base: int + spin_channels: int @property def state_width(self) -> int: """Return the saved-state width: the moments plus both normalizers.""" return self.moment_width + 2 + @property + def has_spin(self) -> bool: + """Return whether the compiled profile carries the spin families.""" + return self.spin_channels > 0 + + +def descriptor_profile( + channels: int, + lmax: int, + has_spin: bool = False, +) -> DescriptorProfile: + """Derive every compiled width from the structural parameters. -def descriptor_profile(channels: int, lmax: int) -> DescriptorProfile: - """Derive every compiled width from the two structural parameters. + The widths mirror ``Profile`` on the device side, + so a mismatch is caught by the operator's own shape validation rather than + producing a silently misread buffer. Parameters ---------- @@ -176,6 +198,9 @@ def descriptor_profile(channels: int, lmax: int) -> DescriptorProfile: Scalar degree-zero width. lmax Maximum angular degree. + has_spin + Whether the native spin families are present. Their width is derived + from the degree profile, so presence is the whole choice. Returns ------- @@ -192,9 +217,32 @@ def descriptor_profile(channels: int, lmax: int) -> DescriptorProfile: bispectrum_dim = int(layout.probe_index.shape[0]) gram_base = degree_channels[0] bispectrum_base = gram_base + gram_total - # Geometric block, the two moment divisors, then the center type tail. + spin_channels = derive_spin_channels(degree_channels) if has_spin else 0 + if has_spin: + # Reduced families, then the node-local on-site vector and quadrupole. + moment_width += 8 * spin_channels + 5 + 8 + # The joint degree-one spin block holds the on-site moment beside the + # isotropic and the bond-projected neighbor channels, so its Gram is + # the upper triangle of a ``1 + 2 C_s`` block. The quadrupole Gram + # drops only its on-site self-term. + vector_width = 1 + 2 * spin_channels + spin_dim = ( + vector_width * (vector_width + 1) // 2 + + 2 + + 2 * degree_channels[2] + + 2 * spin_channels + ) + else: + spin_dim = 0 + # Geometric block, the spin invariants, the two moment divisors, then the + # center type tail. output_width = ( - bispectrum_base + bispectrum_dim + ranks[0] * ranks[1] + 2 + degree_channels[0] + bispectrum_base + + bispectrum_dim + + ranks[0] * ranks[1] + + spin_dim + + 2 + + degree_channels[0] ) return DescriptorProfile( channels=int(channels), @@ -205,6 +253,7 @@ def descriptor_profile(channels: int, lmax: int) -> DescriptorProfile: output_width=output_width, gram_base=gram_base, bispectrum_base=bispectrum_base, + spin_channels=spin_channels, ) @@ -538,14 +587,18 @@ def build_compression_artifacts( ) if not mega_eligible(descriptor): raise ValueError( - "DPA4C compressed CUDA supports channels " - f"{SUPPORTED_CHANNELS} with lmax {SUPPORTED_LMAX} and " + "DPA4C compressed CUDA supports " + f"channels {SUPPORTED_CHANNELS}, lmax {SUPPORTED_LMAX} and " f"radial_modes {SUPPORTED_RADIAL_MODES}, got " f"channels={descriptor.channels}, lmax={descriptor.lmax}, " f"radial_modes={descriptor.radial_modes}" ) device = sample_parameter.device - profile = descriptor_profile(descriptor.channels, descriptor.lmax) + profile = descriptor_profile( + descriptor.channels, + descriptor.lmax, + descriptor.spin is not None, + ) table, info = build_radial_table(descriptor, stride) with torch.no_grad(): @@ -553,8 +606,20 @@ def build_compression_artifacts( device=device, dtype=torch.float32, ) - pair_scale, pair_shift, pair_mixing = descriptor.pair_film.call(type_embedding) + ( + pair_scale, + pair_shift, + pair_mixing, + spin_scale, + spin_shift, + ) = descriptor.pair_film.call(type_embedding) pair_film = torch.stack((pair_scale, pair_shift), dim=-1) + spin_pair, spin_type = _build_spin_caches( + descriptor, + spin_scale, + spin_shift, + device, + ) # The mode axis is innermost so that the coefficients a lane needs for # one channel arrive in one or two vector loads. mixing = ( @@ -578,6 +643,8 @@ def build_compression_artifacts( "info": info, "pair_film": pair_film.detach().contiguous(), "pair_mixing": mixing.detach().contiguous(), + "spin_pair": spin_pair, + "spin_type": spin_type, "type_embedding": type_embedding.detach().contiguous(), "readout_matrices": readout_matrices, "coupling_meta": torch.as_tensor( @@ -600,6 +667,64 @@ def build_compression_artifacts( } +def _build_spin_caches( + descriptor: Any, + spin_scale: torch.Tensor | None, + spin_shift: torch.Tensor | None, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Freeze the two finite tables the native spin branch reads. + + ``spin_pair`` interleaves the ordered scale and shift so that one channel + arrives in a single 64-bit load, matching the geometric PairFiLM cache. + ``spin_type`` packs the four per-type scalars a node needs into one + 128-bit row: the gate divided by the reference magnitude, which conditions + the moment; the bare gate, which the magnetic-coordination family reads + because it counts neighbours that carry a moment rather than the moments + themselves; and the two on-site weights. + + Parameters + ---------- + descriptor + Evaluated pt_expt DPA4C descriptor. + spin_scale, spin_shift + Ordered spin tables, or ``None`` for a spin-free descriptor. + device + Device that receives the packed tables. + + Returns + ------- + spin_pair + Ordered cache with shape ``((T + 1) ** 2, spin_channels, 2)``, or an + empty tensor. + spin_type + Per-type table with shape ``(T + 1, 4)``, or an empty tensor. + """ + empty = torch.zeros(0, dtype=torch.float32, device=device) + if descriptor.spin is None: + return empty, empty + spin = descriptor.spin + gate = spin.spin_mask.to(device=device, dtype=torch.float32) + reference = spin.spin_reference.to(device=device, dtype=torch.float32) + return ( + torch.stack((spin_scale, spin_shift), dim=-1) + .to(device=device, dtype=torch.float32) + .detach() + .contiguous(), + torch.stack( + ( + gate / reference, + gate, + spin.adam_spin_vector_weight.to(device=device, dtype=torch.float32), + spin.adam_spin_quadrupole_weight.to(device=device, dtype=torch.float32), + ), + dim=-1, + ) + .detach() + .contiguous(), + ) + + def _build_readout_matrices( descriptor: Any, profile: DescriptorProfile, @@ -788,15 +913,27 @@ def _cpu_descriptor( rcut: float, eps: float, degree_floor: float, + *, + spin: torch.Tensor | None = None, + spin_pair: torch.Tensor | None = None, + spin_type: torch.Tensor | None = None, ) -> torch.Tensor: - """Reference implementation of the compressed DPA4C descriptor.""" + """Reference implementation of the compressed DPA4C descriptor. + + The native spin block is optional. It comprises the raw per-node moment + ``spin`` with shape ``(N, 3)``, the ordered scale and shift cache + ``spin_pair`` with shape ``((T + 1) ** 2, C_s, 2)``, and the per-type + scalars ``spin_type`` with shape ``(T + 1, 4)``. A spin-free descriptor + omits it and reproduces the geometric descriptor exactly. + """ del destination_order, destination_row_ptr, canonical, coupling_meta compute = edge_vec.to(torch.float32) source, destination = edge_index[0].to(torch.long), edge_index[1].to(torch.long) node_count = atype.shape[0] channels = type_embedding.shape[1] type_count = type_embedding.shape[0] - profile = descriptor_profile(int(channels), int(lmax)) + has_spin = spin is not None and spin.numel() != 0 + profile = descriptor_profile(int(channels), int(lmax), has_spin) radial_modes = 0 if pair_mixing.numel() == 0 else int(pair_mixing.shape[2]) # === Step 1. Build the masked edge geometry === @@ -928,7 +1065,40 @@ def _cpu_descriptor( ) parts[_closed_form_222_coordinate(profile)] = _closed_form_222(tensors) - # === Step 5. Assemble and calibrate the invariant output === + # === Step 5. Reduce and contract the native spin families === + spin_blocks: list[torch.Tensor] = [] + if has_spin: + magnitude, coordination, spin_vector, spin_tensor = _cpu_spin_moments( + spin, + spin_pair, + spin_type, + atype, + source, + destination, + pair_index, + direction, + tabulated, + envelope, + angular_norm, + node_count, + ) + spin_blocks = [ + _half_gram(spin_vector), + # The quadrupole Gram omits its leading diagonal entry, the + # on-site self-term: the harmonic block is homogeneous, so + # |B_2(s)|^2 = |s|^4 makes that entry a per-type constant times + # the square of the vector block's own on-site self-term. + _half_gram(spin_tensor)[:, 1:], + # Cross Gram against the unaligned geometric degree-two moments. + # Both factors carry even spin order, so the product is + # admissible; with a unit direction it evaluates to the + # single-ion anisotropy sum over neighbors. + (spin_tensor.transpose(1, 2) @ blocks[2]).flatten(start_dim=1), + magnitude, + coordination, + ] + + # === Step 6. Assemble and calibrate the invariant output === quartic = (tensor_vector * tensor_vector).sum(dim=-1).flatten(start_dim=1) descriptor = torch.cat( [ @@ -936,6 +1106,7 @@ def _cpu_descriptor( *[_half_gram(block) for block in aligned[1:]], *[parts[key] for key in sorted(parts)], quartic, + *spin_blocks, scalar_divisor, angular_divisor, type_embedding[atype], @@ -945,6 +1116,155 @@ def _cpu_descriptor( return (descriptor - output_mean[None, :]) * output_inv_std[None, :] +def _cpu_spin_moments( + spin: torch.Tensor, + spin_pair: torch.Tensor, + spin_type: torch.Tensor, + atype: torch.Tensor, + source: torch.Tensor, + destination: torch.Tensor, + pair_index: torch.Tensor, + direction: torch.Tensor, + radial: torch.Tensor, + envelope: torch.Tensor, + angular_norm: torch.Tensor, + node_count: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + r"""Reduce the native spin families into their four moment blocks. + + The moment enters through the conditioned node quantity + :math:`\hat{\mathbf s}_i=w_{a_i}\mathbf s_i`, where :math:`w` is the + per-type gate divided by the per-type reference magnitude. Every + neighbor family then shares one edge weight, the spin counterpart of the + geometric amplitude, + + .. math:: + + \phi^{s}_{ij,c}=\chi_{ij}^2\bigl( + \gamma^{s}_{ab,c}g_c(\rho_{ij})+\beta^{s}_{ab,c}\bigr), + + whose squared envelope matches the weight of every non-scalar geometric + moment. The reduced families therefore share the angular normalizer and + need no neighborhood mass of their own. The bond-projected family is the + one that reads the edge direction, and through it the spin branch + contributes to the coordinate gradient angularly as well as radially. + + The on-site channel leads each non-scalar block and is node local: it is + written outside the division, so an invariant pairing it with a neighbor + channel carries exactly one neighborhood normalizer. + + Parameters + ---------- + spin + Raw per-node magnetic moments with shape ``(N, 3)``. + spin_pair + Ordered spin scale and shift with shape ``((T + 1) ** 2, C_s, 2)``. + spin_type + Per-type gate over reference, bare gate, on-site vector weight and + on-site quadrupole weight with shape ``(T + 1, 4)``. + atype + Flat node types with shape ``(N,)``. + source + Source node index of each edge with shape ``(E,)``. + destination + Destination node index of each edge with shape ``(E,)``. + pair_index + Ordered type-pair index of each edge with shape ``(E,)``. + direction + Regularized unit edge directions with shape ``(E, 3)``. + radial + Tabulated distance maps with shape ``(E, channels + radial_modes)``; + the leading ``C_s`` columns are read. + envelope + Masked C³ envelope with shape ``(E,)``. + angular_norm + Reciprocal angular divisor with shape ``(N, 1, 1)``. + node_count + Number of nodes ``N``. + + Returns + ------- + magnitude + Neighbor moment magnitudes with shape ``(N, C_s)``. + coordination + Magnetic effective coordination with shape ``(N, C_s)``. + vector + Joint degree-one spin block with shape ``(N, 3, 1 + 2 C_s)``, holding + the on-site moment, the isotropic neighbor channels and the + bond-projected neighbor channels in that order. + tensor + Degree-two spin block with shape ``(N, 5, 2)``. + """ + spin_channels = int(spin_pair.shape[1]) + weights = spin_type[atype] + conditioned = spin.to(envelope.dtype) * weights[:, 0:1] + neighbor = conditioned[source] + film = spin_pair[pair_index] + weight = (radial[:, :spin_channels] * film[..., 0] + film[..., 1]) * ( + envelope * envelope + )[:, None] + # Component of the neighbor moment along the bond, carried back as a + # vector so that the block Gram turns it into the bond-resolved + # invariants. The masked envelope already zeroes excluded edges, so the + # unmasked direction cannot leak through the amplitude. + bond = direction * (neighbor * direction).sum(dim=-1, keepdim=True) + + # Payload layout: [M0, Mw, V_neighbor, P_neighbor, Q_neighbor] with the + # harmonic component as the outer axis of each non-scalar family, matching + # the geometric moment convention. The magnetic coordination reads the bare + # neighbor gate because it counts neighbors that carry a moment rather + # than the moments themselves, and is therefore nonzero at vanishing spin. + payload = torch.cat( + [ + weight * (neighbor * neighbor).sum(dim=-1, keepdim=True), + weight * weights[source, 1:2], + (neighbor[:, :, None] * weight[:, None, :]).flatten(start_dim=1), + (bond[:, :, None] * weight[:, None, :]).flatten(start_dim=1), + build_angular_basis(neighbor, 2)[:, 4:9] * weight[:, :1], + ], + dim=-1, + ) + reduced = ( + torch.zeros( + node_count, + payload.shape[1], + dtype=payload.dtype, + device=payload.device, + ).index_add_(0, destination, payload) + * angular_norm[:, :, 0] + ) + + onsite_vector = conditioned * weights[:, 2:3] + onsite_tensor = build_angular_basis(conditioned, 2)[:, 4:9] * weights[:, 3:4] + return ( + reduced[:, :spin_channels], + reduced[:, spin_channels : 2 * spin_channels], + torch.cat( + [ + onsite_vector[:, :, None], + reduced[:, 2 * spin_channels : 5 * spin_channels].reshape( + -1, + 3, + spin_channels, + ), + reduced[:, 5 * spin_channels : 8 * spin_channels].reshape( + -1, + 3, + spin_channels, + ), + ], + dim=-1, + ), + torch.cat( + [ + onsite_tensor[:, :, None], + reduced[:, 8 * spin_channels :, None], + ], + dim=-1, + ), + ) + + def _closed_form_112( vectors: torch.Tensor, tensor_vector: torch.Tensor, @@ -999,8 +1319,18 @@ def _closed_form_222_coordinate(profile: DescriptorProfile) -> int: def _cpu_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: """CPU custom-op implementation returning descriptor and opaque state.""" - descriptor = _cpu_descriptor(*args) - profile = descriptor_profile(int(args[9].shape[1]), int(args[17])) + descriptor = _cpu_descriptor( + *args[:16], + *args[19:], + spin=args[16], + spin_pair=args[17], + spin_type=args[18], + ) + profile = descriptor_profile( + int(args[9].shape[1]), + int(args[20]), + args[16].numel() != 0, + ) state = torch.zeros( descriptor.shape[0], profile.state_width, @@ -1027,6 +1357,9 @@ def _forward_fake( coupling_value: torch.Tensor, output_mean: torch.Tensor, output_inv_std: torch.Tensor, + spin: torch.Tensor, + spin_pair: torch.Tensor, + spin_type: torch.Tensor, canonical: bool, lmax: int, table_stride: float, @@ -1049,6 +1382,8 @@ def _forward_fake( coupling_value, output_mean, output_inv_std, + spin_pair, + spin_type, canonical, table_stride, table_max, @@ -1056,7 +1391,9 @@ def _forward_fake( eps, degree_floor, ) - profile = descriptor_profile(int(type_embedding.shape[1]), int(lmax)) + profile = descriptor_profile( + int(type_embedding.shape[1]), int(lmax), spin.numel() != 0 + ) descriptor = torch.empty( atype.shape[0], profile.output_width, @@ -1077,9 +1414,21 @@ def _backward_fake( state: torch.Tensor, edge_vec: torch.Tensor, *args: Any, -) -> torch.Tensor: - del descriptor_gradient, state, args - return torch.empty_like(edge_vec) +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del descriptor_gradient, state + spin = args[15] + has_spin = spin.numel() != 0 + # Each absent output is allocated separately: the schema declares three + # unannotated results, so two of them may not share storage. + return ( + torch.empty_like(edge_vec), + spin.new_empty((spin.shape[0], 3)) + if has_spin + else edge_vec.new_empty((0,), dtype=torch.float32), + torch.empty_like(edge_vec) + if has_spin + else edge_vec.new_empty((0,), dtype=torch.float32), + ) def _cpu_backward( @@ -1087,23 +1436,53 @@ def _cpu_backward( state: torch.Tensor, edge_vec: torch.Tensor, *args: Any, -) -> torch.Tensor: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """CPU custom-op backward returning the coordinate and magnetic cotangents. + + The device operator splits the magnetic cotangent because a + destination-major scan does not own the source of an edge, and the caller + closes it by reducing the per-edge part onto source nodes. The reference + differentiates the whole node axis at once, so it returns the complete + cotangent on the node axis and leaves the per-edge part at zero, which the + same reduction carries through unchanged. + """ del state - if edge_vec.shape[0] == 0: - return torch.zeros_like(edge_vec) + spin = args[15] + has_spin = spin.numel() != 0 value = edge_vec.detach().clone().requires_grad_(True) + moment = spin.detach().clone().requires_grad_(has_spin) with torch.enable_grad(): - descriptor = _cpu_descriptor(value, *args) - (gradient,) = torch.autograd.grad( - (descriptor * descriptor_gradient.to(descriptor.dtype)).sum(), + descriptor = _cpu_descriptor( value, + *args[:15], + *args[18:], + spin=moment, + spin_pair=args[16], + spin_type=args[17], ) - return gradient.to(edge_vec.dtype) + gradients = torch.autograd.grad( + (descriptor * descriptor_gradient.to(descriptor.dtype)).sum(), + (value, moment) if has_spin else (value,), + ) + return ( + gradients[0].to(edge_vec.dtype), + gradients[1] if has_spin else edge_vec.new_empty((0,), dtype=torch.float32), + torch.zeros_like(edge_vec) + if has_spin + else edge_vec.new_empty((0,), dtype=torch.float32), + ) + + +#: Position of the per-node magnetic moment among the operator inputs, and its +#: position among the tensors :func:`_setup_context` saves, which lead with the +#: opaque state. +_SPIN_INPUT_SLOT = 16 +_SPIN_SAVED_SLOT = 1 + _SPIN_INPUT_SLOT def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: - ctx.save_for_backward(output[1], *inputs[:16]) - ctx.scalars = inputs[16:] + ctx.save_for_backward(output[1], *inputs[:19]) + ctx.scalars = inputs[19:] ctx.mark_non_differentiable(output[1]) ctx.set_materialize_grads(False) @@ -1113,15 +1492,35 @@ def _backward( descriptor_gradient: torch.Tensor, state_gradient: torch.Tensor | None, ) -> tuple: + """Return the coordinate cotangent of one compressed descriptor call. + + Raises + ------ + RuntimeError + If the magnetic moment requires a gradient. The operator emits that + cotangent in two pieces and the per-edge piece is reduced onto source + nodes through the source CSR, which the operator schema does not carry; + this registration therefore cannot close the magnetic force and + refuses rather than reporting a vanishing one. + """ del state_gradient tensors = ctx.saved_tensors + if tensors[_SPIN_SAVED_SLOT].requires_grad: + raise RuntimeError( + "deepmd::dpa4c_graph_compress cannot differentiate its magnetic " + "moment through the registered autograd: closing the magnetic " + "force needs the source CSR, which the operator schema does not " + "carry. Call " + "`deepmd.kernels.cuda.dpa4c.graph_compress.dpa4c_graph_compress`, " + "which supplies it." + ) edge_gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( descriptor_gradient, tensors[0], *tensors[1:], *ctx.scalars, - ) - return (edge_gradient,) + (None,) * 22 + )[0] + return (edge_gradient,) + (None,) * 25 _cpu_library: torch.library.Library | None = None @@ -1148,20 +1547,28 @@ def ensure_registered() -> None: ) -def compressed_operator_arguments(descriptor: Any) -> tuple: +def compressed_operator_arguments( + descriptor: Any, + spin: torch.Tensor | None = None, +) -> tuple: """Return the immutable operator arguments of a compressed descriptor. Parameters ---------- descriptor Compressed pt_expt DPA4C descriptor. + spin + Per-node magnetic moments with shape ``(N_all, 3)`` for a + spin-conditioned descriptor. The moment is a runtime input rather than + an artifact; the two tables that condition it are frozen. Returns ------- tuple Radial table, ordered caches, readout projections, coupling tables, - output calibration, and the trailing scalar configuration. + output calibration, and the native spin block. """ + empty = descriptor.compress_spin_type[:0] return ( descriptor.compress_data, descriptor.compress_pair_film, @@ -1173,13 +1580,110 @@ def compressed_operator_arguments(descriptor: Any) -> tuple: descriptor.compress_coupling_value, descriptor.compress_output_mean, descriptor.compress_output_inv_std, + empty if spin is None else spin.to(torch.float32).contiguous(), + descriptor.compress_spin_pair, + descriptor.compress_spin_type, ) +def reduce_edge_spin_gradient( + edge_spin_gradient: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, +) -> torch.Tensor: + """Reduce a per-edge magnetic cotangent onto its source nodes. + + The source CSR groups the outgoing edges of a node contiguously, so the + reduction is a segment sum over a gathered edge axis. That fixes the + summation order from the topology rather than from arrival order, which an + atomic scatter would not. + + The permutation is gathered in full rather than sliced to the physical + edge count. The segments consume only the leading rows, which the source + grouping already reserves for the physical edges, and reading the count + off the row pointers would turn a device value into a Python integer that + symbolic tracing cannot resolve. + + Parameters + ---------- + edge_spin_gradient + Per-edge magnetic cotangent with shape ``(E, 3)``. + source_order + Source-grouped edge permutation with shape ``(E,)``. + source_row_ptr + Source CSR offsets with shape ``(N + 1,)``. + + Returns + ------- + torch.Tensor + Per-node magnetic cotangent with shape ``(N, 3)``. + """ + ordered = torch.index_select( + edge_spin_gradient, + 0, + source_order.to(torch.int64), + ) + return torch.segment_reduce( + ordered, + "sum", + lengths=(source_row_ptr[1:] - source_row_ptr[:-1]).to(torch.int64), + axis=0, + unsafe=True, + ) + + +class _CompressedDescriptor(torch.autograd.Function): + """Autograd wrapper that closes the magnetic force of the level-one path. + + The operator emits the on-site magnetic gradient per node and the + neighbour part per edge, because a destination-major scan does not own the + source of an edge. Reducing the second onto source nodes needs the source + CSR, which the operator schema does not carry, so the reduction lives here + where the graph is in scope. + """ + + @staticmethod + def forward( + ctx: Any, + edge_vec: torch.Tensor, + spin: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + operator_args: tuple, + ) -> torch.Tensor: + descriptor, state = torch.ops.deepmd.dpa4c_graph_compress( + edge_vec, *operator_args + ) + ctx.save_for_backward(state, edge_vec, source_order, source_row_ptr) + ctx.operator_args = operator_args + return descriptor + + @staticmethod + def backward(ctx: Any, descriptor_gradient: torch.Tensor) -> tuple: + state, edge_vec, source_order, source_row_ptr = ctx.saved_tensors + ( + edge_gradient, + spin_gradient, + edge_spin_gradient, + ) = torch.ops.deepmd.dpa4c_graph_compress_backward( + descriptor_gradient.contiguous(), + state, + edge_vec, + *ctx.operator_args, + ) + spin_gradient = spin_gradient + reduce_edge_spin_gradient( + edge_spin_gradient, + source_order, + source_row_ptr, + ) + return edge_gradient, spin_gradient, None, None, None + + def dpa4c_graph_compress( descriptor: Any, graph: Any, atype: torch.Tensor, + spin: torch.Tensor | None = None, ) -> torch.Tensor: """Evaluate the compressed DPA4C graph descriptor. @@ -1191,6 +1695,8 @@ def dpa4c_graph_compress( NeighborGraph with destination CSR topology. atype Flat node types with shape ``(N,)``. + spin + Per-node magnetic moments with shape ``(N, 3)``, or ``None``. Returns ------- @@ -1206,18 +1712,35 @@ def dpa4c_graph_compress( ensure_registered() if graph.destination_order is None or graph.destination_row_ptr is None: raise ValueError("DPA4C compressed CUDA requires destination CSR topology") - descriptor_output, _state = torch.ops.deepmd.dpa4c_graph_compress( - graph.edge_vec.contiguous(), + operator_args = ( graph.edge_index.contiguous(), graph.edge_mask.contiguous(), graph.destination_order.contiguous(), graph.destination_row_ptr.contiguous(), atype.contiguous(), - *compressed_operator_arguments(descriptor), + *compressed_operator_arguments(descriptor, spin), bool(graph.destination_sorted), int(descriptor.lmax), *descriptor._compression_scalars, ) + if spin is not None and spin.requires_grad: + if graph.source_order is None or graph.source_row_ptr is None: + raise ValueError( + "DPA4C compressed CUDA requires source CSR topology to close " + "the magnetic force" + ) + descriptor_output = _CompressedDescriptor.apply( + graph.edge_vec.contiguous(), + spin, + graph.source_order.contiguous(), + graph.source_row_ptr.contiguous(), + operator_args, + ) + else: + descriptor_output, _state = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec.contiguous(), + *operator_args, + ) return descriptor_output.to(graph.edge_vec.dtype) @@ -1230,8 +1753,16 @@ def dpa4c_graph_compress_energy_force( atom_bias: torch.Tensor, node_capacity: int, do_atomic_virial: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Evaluate compressed DPA4C energy, force, and virial without a tape. + spin: torch.Tensor | None = None, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Evaluate compressed DPA4C energy, force, virial and magnetic force. Parameters ---------- @@ -1251,6 +1782,8 @@ def dpa4c_graph_compress_energy_force( Force-scatter node capacity. do_atomic_virial Whether to return per-node virials. + spin + Per-node magnetic moments with shape ``(N, 3)``, or ``None``. Returns ------- @@ -1264,6 +1797,8 @@ def dpa4c_graph_compress_energy_force( Per-frame virial with shape ``(F, 3, 3)``, fp32. atom_virial Per-node virial with shape ``(N, 3, 3)`` or an empty tensor. + force_mag + Per-node magnetic force with shape ``(N, 3)``, or an empty tensor. Raises ------ @@ -1300,7 +1835,7 @@ def dpa4c_graph_compress_energy_force( graph.destination_order.contiguous(), graph.destination_row_ptr.contiguous(), atype.contiguous(), - *compressed_operator_arguments(descriptor), + *compressed_operator_arguments(descriptor, spin), bool(graph.destination_sorted), int(descriptor.lmax), *descriptor._compression_scalars, @@ -1320,13 +1855,20 @@ def dpa4c_graph_compress_energy_force( graph.n_node, ) del node_descriptor - edge_gradient = torch.ops.deepmd.dpa4c_graph_compress_backward( + ( + edge_gradient, + spin_gradient, + edge_spin_gradient, + ) = torch.ops.deepmd.dpa4c_graph_compress_backward( descriptor_gradient, state, edge_vec, *operator_args, ) - force, atom_virial, virial = edge_force_virial( + # The on-site magnetic gradient closes in the node kernel; the neighbour + # part belongs to source nodes, and the force assembly already walks that + # grouping, so it is reduced there rather than in a pass of its own. + force, atom_virial, virial, source_spin = edge_force_virial( edge_gradient, edge_vec, graph.edge_index, @@ -1336,10 +1878,12 @@ def dpa4c_graph_compress_energy_force( graph.source_order, graph.source_row_ptr, graph.n_node, + edge_spin_gradient, node_capacity, do_atomic_virial, ) - return energy, atom_energy, force, virial, atom_virial + force_mag = spin_gradient if spin is None else -(spin_gradient + source_spin) + return energy, atom_energy, force, virial, atom_virial, force_mag def fitting_energy_and_gradient( diff --git a/deepmd/kernels/cuda/edge_force_virial.py b/deepmd/kernels/cuda/edge_force_virial.py index e7ea7b772d..eb254bb605 100644 --- a/deepmd/kernels/cuda/edge_force_virial.py +++ b/deepmd/kernels/cuda/edge_force_virial.py @@ -121,14 +121,16 @@ def _fake( source_order: torch.Tensor, source_row_ptr: torch.Tensor, n_node_per_frame: torch.Tensor, + edge_spin_gradient: torch.Tensor, node_capacity: int, want_atom_virial: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: n_frame = n_node_per_frame.shape[0] return ( g_e.new_empty(node_capacity, 3), g_e.new_empty(node_capacity if want_atom_virial else 0, 3, 3), g_e.new_empty(n_frame, 3, 3), + g_e.new_empty(node_capacity if edge_spin_gradient.numel() else 0, 3), ) @@ -139,15 +141,17 @@ def _canonical_fake( source_row_ptr: torch.Tensor, source_order: torch.Tensor, n_node_per_frame: torch.Tensor, + edge_spin_gradient: torch.Tensor, node_capacity: int, want_atom_virial: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: del edge_vec, destination_row_ptr, source_row_ptr, source_order n_frame = n_node_per_frame.shape[0] return ( g_e.new_empty(node_capacity, 3), g_e.new_empty(node_capacity if want_atom_virial else 0, 3, 3), g_e.new_empty(n_frame, 3, 3), + g_e.new_empty(node_capacity if edge_spin_gradient.numel() else 0, 3), ) @@ -161,9 +165,10 @@ def _cpu( source_order: torch.Tensor, source_row_ptr: torch.Tensor, n_node_per_frame: torch.Tensor, + edge_spin_gradient: torch.Tensor, node_capacity: int, want_atom_virial: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: from deepmd.dpmodel.utils.neighbor_graph import edge_force_virial as reference force, atom_virial, virial = reference( @@ -176,7 +181,18 @@ def _cpu( ) if not want_atom_virial: atom_virial = atom_virial.new_zeros(0, 3, 3) - return force, atom_virial, virial + if edge_spin_gradient.numel(): + # A masked edge carries no force and no moment, so the two reductions + # must agree on which edges exist. + contribution = edge_spin_gradient + if edge_mask.numel(): + contribution = contribution * edge_mask[:, None].to(contribution.dtype) + magnetic_force = torch.zeros( + node_capacity, 3, dtype=g_e.dtype, device=g_e.device + ).index_add_(0, edge_index[0], contribution) + else: + magnetic_force = g_e.new_zeros(0, 3) + return force, atom_virial, virial, magnetic_force def _canonical_cpu( @@ -186,9 +202,10 @@ def _canonical_cpu( source_row_ptr: torch.Tensor, source_order: torch.Tensor, n_node_per_frame: torch.Tensor, + edge_spin_gradient: torch.Tensor, node_capacity: int, want_atom_virial: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: physical_edge_count = int(destination_row_ptr[-1].item()) node_count = destination_row_ptr.shape[0] - 1 destination = torch.repeat_interleave( @@ -232,6 +249,7 @@ def _canonical_cpu( source_order, source_row_ptr, n_node_per_frame, + edge_spin_gradient, node_capacity, want_atom_virial, ) @@ -277,9 +295,10 @@ def edge_force_virial( source_order: torch.Tensor, source_row_ptr: torch.Tensor, n_node_per_frame: torch.Tensor, + edge_spin_gradient: torch.Tensor, node_capacity: int, want_atom_virial: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Assemble force and virial from the per-edge energy gradient. Matches the array-API reference up to floating summation order: @@ -304,6 +323,9 @@ def edge_force_virial( Destination/source CSR offsets with shape (N + 1,), int64. n_node_per_frame : torch.Tensor Per-frame node counts with shape (nf,), int64. + edge_spin_gradient : torch.Tensor + Per-edge magnetic cotangent with shape (E, 3), or an empty tensor when + the model carries no magnetic degree of freedom. node_capacity : int Padded node-axis size ``N`` (may be a ``SymInt`` under tracing). want_atom_virial : bool @@ -318,6 +340,9 @@ def edge_force_virial( when not requested. virial : torch.Tensor Per-frame virial with shape (nf, 3, 3). + magnetic_force : torch.Tensor + Per-source total of the magnetic cotangent with shape (N, 3), or an + empty (0, 3) tensor when no spin cotangent was supplied. """ ensure_registered() return torch.ops.deepmd.edge_force_virial( @@ -330,6 +355,7 @@ def edge_force_virial( source_order.contiguous(), source_row_ptr.contiguous(), n_node_per_frame, + edge_spin_gradient.contiguous(), node_capacity, want_atom_virial, ) @@ -342,9 +368,10 @@ def canonical_edge_force_virial( source_row_ptr: torch.Tensor, source_order: torch.Tensor, n_node_per_frame: torch.Tensor, + edge_spin_gradient: torch.Tensor, node_capacity: int, want_atom_virial: bool, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Assemble force and virial from a compact canonical edge stream. Parameters @@ -359,6 +386,8 @@ def canonical_edge_force_virial( Edge storage positions grouped by source with shape ``(S,)``. n_node_per_frame Per-frame node counts with shape ``(nf,)``. + edge_spin_gradient + Per-edge magnetic cotangent with shape ``(E, 3)``, or empty. node_capacity Flat node count ``N``. want_atom_virial @@ -367,7 +396,9 @@ def canonical_edge_force_virial( Returns ------- tuple[torch.Tensor, torch.Tensor, torch.Tensor] - Force, optional atom virial, and frame virial. + Force, optional atom virial, frame virial, and the per-source total of + the magnetic cotangent, the last empty when no spin cotangent was + supplied. """ ensure_registered() return torch.ops.deepmd.canonical_edge_force_virial( @@ -377,6 +408,7 @@ def canonical_edge_force_virial( source_row_ptr.contiguous(), source_order.contiguous(), n_node_per_frame, + edge_spin_gradient.contiguous(), node_capacity, want_atom_virial, ) diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index 91266b8788..cb1172c369 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -309,6 +309,29 @@ def _type_pair_table(desc: Any, type_embedding: torch.Tensor) -> torch.Tensor: return se.cal_g_strip(two_side, 0) +def _without_magnetic_force( + output: tuple[torch.Tensor, ...], +) -> tuple[torch.Tensor, ...]: + """Append the empty magnetic force of a spin-free fused composition. + + Every implementation of ``fused_energy_force_graph`` returns the same six + outputs, so the caller reads the magnetic force by position rather than by + probing the arity of the result. + + Parameters + ---------- + output + The five spin-free outputs: energy, atom energy, force, virial and + atom virial. + + Returns + ------- + tuple of torch.Tensor + The same outputs followed by an empty magnetic force. + """ + return (*output, output[2].new_empty((0, 3))) + + @BaseDescriptor.register("se_atten") @BaseDescriptor.register("dpa1") @torch_module @@ -1068,6 +1091,7 @@ def fused_energy_force_graph( ownership: torch.Tensor, atom_bias: torch.Tensor, do_atomic_virial: bool, + spin: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...] | None: """End-to-end fused energy / force / virial from the edge stream. @@ -1123,7 +1147,28 @@ def fused_energy_force_graph( if not (ef_op_available() and mega_eligible(self)): return None - return dpa1_graph_compress_energy_force( + return _without_magnetic_force( + dpa1_graph_compress_energy_force( + self, + fit, + graph, + atype, + type_embedding, + ownership, + atom_bias, + node_capacity=node_capacity, + do_atomic_virial=do_atomic_virial, + ) + ) + from deepmd.kernels.cuda.dpa1.graph_energy_force import ( + dpa1_graph_energy_force, + op_available, + ) + + if not op_available(): + return None + return _without_magnetic_force( + dpa1_graph_energy_force( self, fit, graph, @@ -1134,23 +1179,6 @@ def fused_energy_force_graph( node_capacity=node_capacity, do_atomic_virial=do_atomic_virial, ) - from deepmd.kernels.cuda.dpa1.graph_energy_force import ( - dpa1_graph_energy_force, - op_available, - ) - - if not op_available(): - return None - return dpa1_graph_energy_force( - self, - fit, - graph, - atype, - type_embedding, - ownership, - atom_bias, - node_capacity=node_capacity, - do_atomic_virial=do_atomic_virial, ) def _call_graph_triton( diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py index 21b3048c06..1a9fd300b8 100644 --- a/deepmd/pt_expt/descriptor/dpa4c.py +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -34,6 +34,14 @@ _TRAINABLE_ATTRS: dict[str, tuple[str, ...]] = { "SeZMTypeEmbedding": ("adam_type_embedding",), "RadialBasis": ("adam_freqs",), + "OrderedPairFiLM": ( + "adam_spin_scale_anchor", + "adam_spin_shift_anchor", + ), + "SpinChannels": ( + "adam_spin_vector_weight", + "adam_spin_quadrupole_weight", + ), } @@ -124,6 +132,7 @@ def call_graph( atype: torch.Tensor, type_embedding: torch.Tensor | None = None, comm_dict: dict | None = None, + spin: torch.Tensor | None = None, ) -> tuple[torch.Tensor, None]: """Evaluate the graph descriptor with compressed CUDA dispatch. @@ -137,6 +146,9 @@ def call_graph( Optional complete DPA4 type table. comm_dict Communication metadata accepted by the common graph ABI; unused. + spin + Per-node spin with shape ``(N, 3)``, mandatory for a + spin-conditioned descriptor. Returns ------- @@ -160,10 +172,14 @@ def call_graph( ) if op_available() and mega_eligible(self): + # The operator conditions the moment on device from its frozen + # per-type table, so it takes the raw input rather than the + # output of ``SpinChannels.conditioned_spin``. return dpa4c_graph_compress( self, graph, atype, + None if self.spin is None else self.require_spin(spin), ), None if type_embedding is None: type_embedding = self.type_embedding.call() @@ -172,13 +188,14 @@ def call_graph( atype, type_embedding=type_embedding, comm_dict=comm_dict, + spin=spin, ) def build_edge_features( self, graph: Any, *args: Any, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: """Build the edge features under the DPA4C mixed-precision policy. The per-edge stage is the only region DPA4C autocasts. It holds every @@ -198,13 +215,20 @@ def build_edge_features( The two are independent: mixed precision at inference is a throughput choice that must not require a model to have been trained with it. + A spin-conditioned descriptor never autocasts. Its scalar and + quadrupole families are quadratic in the magnetic moment and feed a + fourth-order readout, and the magnetic force differentiates them + twice, so the eight mantissa bits of bfloat16 are not an acceptable + trade for a configuration whose throughput is not the binding + constraint. + Parameters ---------- graph Neighbor graph in descriptor compute precision. *args - Node types and ordered pair tables forwarded unchanged to the - backend-neutral implementation. + Node types, the ordered pair cache, and the conditioned spin, + forwarded unchanged to the backend-neutral implementation. Returns ------- @@ -214,16 +238,22 @@ def build_edge_features( Masked Cartesian harmonics with shape ``(E, (lmax + 1) ** 2)``. envelope Masked C³ envelope with shape ``(E,)``. + spin_payload + Masked per-edge spin payload, or ``None``. """ - autocast = graph.edge_vec.device.type == "cuda" and ( - self.use_amp if self.training else self.use_amp_infer + autocast = ( + self.spin is None + and graph.edge_vec.device.type == "cuda" + and (self.use_amp if self.training else self.use_amp_infer) ) if not autocast: return super().build_edge_features(graph, *args) with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): features = super().build_edge_features(graph, *args) dtype = graph.edge_vec.dtype - return tuple(feature.to(dtype) for feature in features) + return tuple( + None if feature is None else feature.to(dtype) for feature in features + ) def _apply_autocast_policy(self) -> None: """Let the layers inside the autocast region emit reduced precision. @@ -232,10 +262,10 @@ def _apply_autocast_policy(self) -> None: which would undo autocast at every layer of the radial network. Only the radial trunk and the mode head sit inside the region, so only they are opted out of that restoration, and only when mixed precision can - actually engage. A descriptor with neither switch set keeps the - default behavior exactly. + actually engage. A descriptor with neither switch set, and every + spin-conditioned descriptor, keeps the default behavior exactly. """ - enabled = self.use_amp or self.use_amp_infer + enabled = self.spin is None and (self.use_amp or self.use_amp_infer) for layer in self.radial_embedding.layers: layer.autocast_output = enabled if self.radial_mode_head is not None: @@ -426,6 +456,18 @@ def train(self, mode: bool = True) -> "DescrptDPA4C": ) return super().train(mode) + def compression_needs_min_nbor_dist(self) -> bool: + """Return whether compression consumes the minimum neighbor distance. + + Returns + ------- + bool + Always ``False``. The radial table spans ``[0, rcut]``, a domain + fixed by the cutoff rather than by the training data, so the + caller can skip the neighbor-statistics pass. + """ + return False + def enable_compression( self, min_nbor_dist: float, @@ -476,6 +518,7 @@ def fused_energy_force_graph( ownership: torch.Tensor, atom_bias: torch.Tensor, do_atomic_virial: bool, + spin: torch.Tensor | None = None, ) -> ( tuple[ torch.Tensor, @@ -483,13 +526,15 @@ def fused_energy_force_graph( torch.Tensor, torch.Tensor, torch.Tensor, + torch.Tensor, ] | None ): """Evaluate the inference-only compressed energy-force composition. Returns ``None`` when the model or graph cannot use the level-two CUDA - path, allowing the caller to retain the generic autograd lower. + path, allowing the caller to retain the generic autograd lower. The + trailing output is the magnetic force, empty for a spin-free model. """ if ( self.training @@ -526,4 +571,5 @@ def fused_energy_force_graph( atom_bias, atype.shape[0], do_atomic_virial, + None if self.spin is None else self.require_spin(spin), ) diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 3d9bb6506a..9aeae29ca3 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -45,10 +45,18 @@ def enable_compression( model_dict = serialize_from_file(input_file) model = BaseModel.deserialize(model_dict["model"]) - # 2. Get or compute min_nbor_dist + # 2. Get or compute min_nbor_dist. Measuring it is a dense all-pairs pass + # over the training data, so it is only run for models that tabulate + # from the shortest observed distance. min_nbor_dist = model.get_min_nbor_dist() if min_nbor_dist is None: min_nbor_dist = model_dict.get("min_nbor_dist") + if min_nbor_dist is None and not model.compression_needs_min_nbor_dist(): + log.info( + "The model tabulates over an analytically bounded domain; " + "skipping the neighbor statistics." + ) + min_nbor_dist = 0.0 if min_nbor_dist is None: log.info( "Minimal neighbor distance is not saved in the model, " diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index db32f0ceec..69f388b5da 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1823,7 +1823,7 @@ def _eval_model_spin( request_defs: list[OutputVariableDef], charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: - if self.metadata.get("lower_input_kind") == "graph": + if self.metadata.get("lower_input_kind") in ("graph", "dpa4c_canonical"): # Native-spin (NeighborGraph route): no virtual atoms and no # extended/nlist ABI at all -- dispatch to the graph-native fast # path (mirrors _eval_model's dispatch to _eval_model_graph for @@ -2102,22 +2102,70 @@ def _eval_model_graph_spin( # the same axis as ``atype``/``spin`` (mirrors _eval_model_graph). aparam_t = aparam_t.reshape(nframes * natoms, -1) - model_inputs = ( - atype_t, - n_node_t, - n_node_t, - edge_index_t, - edge_vec_t, - edge_mask_t, - destination_order_t, - destination_row_ptr_t, - source_order_t, - source_row_ptr_t, - spin_t, - fparam_t, - aparam_t, - self._make_charge_spin_input(nframes, charge_spin), - ) + if self.metadata.get("lower_input_kind") == "dpa4c_canonical": + # The compact canonical ABI has no conditional tail, so the moment + # is the last slot rather than the eleventh, and any conditioning + # input the artifact cannot receive is an error rather than a + # silently dropped argument. + if ( + self.get_dim_fparam() > 0 + or self.get_dim_aparam() > 0 + or int(self.metadata.get("dim_chg_spin", 0) or 0) > 0 + ): + raise NotImplementedError( + "compact canonical artifacts carry no fparam/aparam/" + "charge_spin inputs; a model requiring them must not be " + "frozen with a canonical lower kind." + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + ) + + compact = canonical_graph_from_neighbor_graph( + NeighborGraph( + n_node=n_node_t, + edge_index=edge_index_t, + edge_vec=edge_vec_t, + edge_mask=edge_mask_t, + n_local=n_node_t, + destination_order=destination_order_t, + destination_row_ptr=destination_row_ptr_t, + source_order=source_order_t, + source_row_ptr=source_row_ptr_t, + destination_sorted=bool(graph.destination_sorted), + ) + ) + model_inputs = ( + atype_t, + compact.n_node, + compact.n_local, + compact.source, + compact.edge_vec, + compact.destination_row_ptr, + compact.source_row_ptr, + compact.source_order, + spin_t.to(torch.float32), + ) + else: + model_inputs = ( + atype_t, + n_node_t, + n_node_t, + edge_index_t, + edge_vec_t, + edge_mask_t, + destination_order_t, + destination_row_ptr_t, + source_order_t, + source_row_ptr_t, + spin_t, + fparam_t, + aparam_t, + self._make_charge_spin_input(nframes, charge_spin), + ) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 7838781889..7aaa9cb8dd 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -129,7 +129,7 @@ def edge_energy_deriv( and source_row_ptr is not None ): n_cap = node_capacity if node_capacity is not None else int(n_node.sum()) - force, atom_virial, virial = fused_edge_force_virial( + force, atom_virial, virial, _ = fused_edge_force_virial( g_e, edge_vec, edge_index, @@ -139,6 +139,7 @@ def edge_energy_deriv( source_order, source_row_ptr, n_node, + edge_vec.new_zeros(0, 3), n_cap, do_atomic_virial, ) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 90c85cfd5c..b053d0f913 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -66,6 +66,7 @@ def forward_lower_canonical_graph( source_order: torch.Tensor, *, do_atomic_virial: bool, + spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Evaluate an eligible compressed canonical deployment graph. @@ -90,6 +91,9 @@ def forward_lower_canonical_graph( dtype as ``source``. do_atomic_virial Whether to return the per-node virial. + spin + Per-node magnetic moments with shape ``(N, 3)`` for a native-spin + model, or ``None``. Ghost rows carry their owner's moment. Returns ------- @@ -136,7 +140,7 @@ def forward_lower_canonical_graph( fitting = self.atomic_model.fitting_net atom_bias = fitting.bias_atom_e[:, 0] + self.atomic_model.out_bias[0, :, 0] if use_dpa4c: - energy, atom_energy, force, virial, atom_virial = ( + energy, atom_energy, force, virial, atom_virial, force_mag = ( dpa4c_canonical_compress_energy_force( descriptor, fitting, @@ -145,9 +149,11 @@ def forward_lower_canonical_graph( output_mask, atom_bias, do_atomic_virial, + spin, ) ) else: + force_mag = None energy, atom_energy, force, virial, atom_virial = ( dpa1_canonical_compress_energy_force( descriptor, @@ -169,6 +175,8 @@ def forward_lower_canonical_graph( "virial": virial, "mask": output_mask.to(torch.int32), } + if force_mag is not None and force_mag.numel() != 0: + result["force_mag"] = force_mag if do_atomic_virial: result["atom_virial"] = atom_virial return result diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 7aa82215b2..563bd4f1a4 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -87,6 +87,7 @@ def _fused_energy_force_graph( graph: Any, atype: torch.Tensor, do_atomic_virial: bool, + spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor] | None: """End-to-end energy / force / virial via fused inference operators. @@ -115,10 +116,13 @@ def _fused_energy_force_graph( output_mask, atom_bias, do_atomic_virial, + spin, ) if out is None: return None - energy, atom_energy, force, virial, atom_virial = out + # Every implementation returns the same six outputs; a descriptor without + # native spin leaves the magnetic force empty. + energy, atom_energy, force, virial, atom_virial, force_mag = out n = atype.shape[0] nf = graph.n_node.shape[0] var = fit.var_name @@ -129,6 +133,12 @@ def _fused_energy_force_graph( var + "_derv_c_redu": virial.reshape(nf, 1, 9), "mask": output_mask.to(torch.int32), } + if force_mag.numel() != 0: + ret[var + "_derv_r_mag"] = force_mag.reshape(n, 1, 3) + elif spin is not None: + # A moment was supplied but this descriptor emits no magnetic force, + # so the fused result is incomplete and the caller must fall back. + return None if do_atomic_virial: ret[var + "_derv_c"] = atom_virial.reshape(n, 1, 9) return ret @@ -675,10 +685,13 @@ def forward_common_lower_graph( ) # Level 2 emits force as a value through the inference-only custom # operator pipeline. Ineligible models use the autograd lower. - # The fused pipeline has no mag output, so spin-conditioned models - # always take the autograd lower below. - if not self.training and cuda_infer_level() >= 2 and spin is None: - fused = _fused_energy_force_graph(self, graph, atype, do_atomic_virial) + # The fused pipeline emits the magnetic force as a value for a + # descriptor that declares native spin, and returns nothing when it + # cannot serve the request at all. + if not self.training and cuda_infer_level() >= 2: + fused = _fused_energy_force_graph( + self, graph, atype, do_atomic_virial, spin + ) if fused is not None: return fused atomic_ret = self.atomic_model.forward_common_atomic_graph( diff --git a/deepmd/pt_expt/model/native_spin_model.py b/deepmd/pt_expt/model/native_spin_model.py index bba9d104b3..5103af1d03 100644 --- a/deepmd/pt_expt/model/native_spin_model.py +++ b/deepmd/pt_expt/model/native_spin_model.py @@ -305,6 +305,94 @@ def fn( charge_spin, ) + def forward_lower_canonical_graph_exportable( + self, + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + source: torch.Tensor, + edge_vec: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_row_ptr: torch.Tensor, + source_order: torch.Tensor, + spin: torch.Tensor, + *, + do_atomic_virial: bool, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + """Trace the compact canonical spin lower into an exportable module. + + THIS METHOD OWNS the positional ``.pt2`` ABI for compact canonical + spin models (mirrored verbatim by the C++ / serialization seams): + ``spin`` sits at index 8, directly after the dual-CSR block, which is + the same placement rule the graph ABI uses at index 10. The compact ABI + has no conditional tail, so index 8 is the last slot. + + The magnetic force is a value here rather than an autograd result: the + fused operator emits the on-site half from the node kernel and the + neighbour half per edge, and the deployment path reduces the second + onto source nodes. ``mask_mag`` is derived from the types by the same + single owner the eager translation uses. + + Parameters + ---------- + atype, n_node, n_local, source, edge_vec + Compact graph node and edge tensors. + destination_row_ptr, source_row_ptr, source_order + Compact dual-CSR topology. + spin + ``(N, 3)`` per-node moment. Ghost rows carry their owner's moment. + do_atomic_virial + Whether the traced output includes per-node virial. + **make_fx_kwargs + Additional arguments passed to :func:`make_fx`. + + Returns + ------- + torch.nn.Module + Traced nine-input compact deployment module whose output dict adds + ``force_mag`` and ``mask_mag`` to the spin-free key set. + """ + model = self + + def fn( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + source: torch.Tensor, + edge_vec: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_row_ptr: torch.Tensor, + source_order: torch.Tensor, + spin: torch.Tensor, + ) -> dict[str, torch.Tensor]: + result = model.forward_lower_canonical_graph( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + do_atomic_virial=do_atomic_virial, + spin=spin, + ) + result["mask_mag"] = model._spin_active_mask(atype) + return result + + return make_fx(fn, **make_fx_kwargs)( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + spin, + ) + def forward_lower_graph_exportable_with_comm( self, atype: torch.Tensor, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 18cafbff4e..82193ba339 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1496,6 +1496,13 @@ def _forward_graph( ``(N, 3)`` with ``N == nframes * nloc`` for a single-rank carry-all graph, so no extended->local scatter is needed; only the flat ``(N, *)`` node keys are unravelled to ``(nf, nloc, *)`` at the I/O boundary. + + A native-spin model additionally passes the per-node moment, which the + compiled lower turns into ``force_mag`` (and the type gate ``mask_mag``) + alongside the energy keys. Presence of the moment is a property of the + model, not of the batch, and it is part of the structure key + (:func:`_get_model_structure_key`), so a cached graph is never shared + between a spin and a spin-free task. """ from deepmd.dpmodel.utils.neighbor_graph import ( compact_nodes, @@ -1610,6 +1617,8 @@ def _forward_graph( # Feed a detached, grad-enabled edge_vec leaf: the traced graph's internal # ``edge_vec.detach()`` is stripped by ``_strip_saved_tensor_detach`` (as # for the dense ext_coord leaf), so the force backward roots at this input. + # ``spin`` is the graph lower's second leaf and gets the same treatment, + # rooting the magnetic-force backward at the moment input. edge_vec = ng.edge_vec.detach().requires_grad_(True) if spin is not None: spin = spin.detach().requires_grad_(True) @@ -2364,17 +2373,10 @@ def _raise_if_full_validation_unsupported( "training with training.zero_stage < 2." ) - if self.models[DEFAULT_TASK_KEY].has_spin() or isinstance( - self.loss, EnergySpinLoss - ): + if not isinstance(self.loss, (EnergyLoss, EnergySpinLoss)): raise ValueError( "validating.full_validation only supports single-task energy " - "training; spin-energy training is not supported." - ) - - if not isinstance(self.loss, EnergyLoss): - raise ValueError( - "validating.full_validation only supports single-task energy training." + "or spin-energy training." ) if validation_data is None: diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py index 32af3a0c42..1af9e454bc 100644 --- a/deepmd/pt_expt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -88,7 +88,7 @@ BEST_CKPT_PREFIX = "best.ckpt" EMA_BEST_CKPT_PREFIX = "best_ema.ckpt" VAL_LOG_SIGNIFICANT_DIGITS = 5 -VAL_LOG_COLUMN_GAP = " " +VAL_LOG_COLUMN_GAP = " " VAL_LOG_HEADER_PREFIX = "# " VAL_LOG_DATA_PREFIX = " " @@ -285,7 +285,7 @@ def __init__( ) header_label = f"{column_name}({metric_unit})" self.table_column_specs.append( - (metric_key, header_label, max(len(header_label), 18)) + (metric_key, header_label, max(len(header_label), 10)) ) self.topk_records = self._load_topk_records() @@ -525,9 +525,7 @@ def _evaluate_system( natoms = int(test_data["type"].shape[1]) nframes = int(test_data["coord"].shape[0]) include_virial = ( - not self.profile.needs_spin - and data_system.pbc - and bool(test_data.get("find_virial", 0.0)) + data_system.pbc and bool(test_data.get("find_virial", 0.0)) ) spin = ( test_data["spin"].reshape(nframes, -1) if self.profile.needs_spin else None diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 4abc487ad3..94250f36fa 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -614,8 +614,14 @@ def build_synthetic_canonical_graph_inputs( e_max: int, *, device: torch.device, + want_spin: bool = False, ) -> tuple[torch.Tensor, ...]: - """Build compact canonical trace inputs for compressed CUDA descriptors.""" + """Build compact canonical trace inputs for compressed CUDA descriptors. + + ``want_spin`` appends the per-node moment at slot 8, the last slot of the + compact ABI, matching + :meth:`~deepmd.pt_expt.model.native_spin_model.NativeSpinEnergyModel.forward_lower_canonical_graph_exportable`. + """ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) @@ -632,6 +638,7 @@ def build_synthetic_canonical_graph_inputs( want_fparam=False, want_aparam=False, want_charge_spin=False, + want_spin=want_spin, ) ( atype, @@ -644,10 +651,9 @@ def build_synthetic_canonical_graph_inputs( destination_row_ptr, source_order, source_row_ptr, - _fparam, - _aparam, - _charge_spin, + *tail, ) = sample + spin = tail[0] if want_spin else None graph = NeighborGraph( n_node=n_node, edge_index=edge_index, @@ -661,7 +667,7 @@ def build_synthetic_canonical_graph_inputs( destination_sorted=True, ) compact = canonical_graph_from_neighbor_graph(graph) - return ( + compact_sample = ( atype, compact.n_node, compact.n_local, @@ -671,13 +677,17 @@ def build_synthetic_canonical_graph_inputs( compact.source_row_ptr, compact.source_order, ) + return compact_sample if spin is None else (*compact_sample, spin) def _build_canonical_graph_dynamic_shapes( *sample_inputs: torch.Tensor, ) -> tuple: - """Build dynamic shapes for the eight-tensor compact deployment ABI.""" - del sample_inputs + """Build dynamic shapes for the compact deployment ABI. + + The trailing spin slot is present only for a native-spin model, so the + sample length selects the shape tuple. + """ from deepmd.pt_expt.utils.canonical_graph import ( UINT32_MAX, ) @@ -689,7 +699,7 @@ def _build_canonical_graph_dynamic_shapes( min=2, max=UINT32_MAX, ) - return ( + shapes = ( {0: node_dim}, {0: nframes_dim}, {0: nframes_dim}, @@ -699,6 +709,7 @@ def _build_canonical_graph_dynamic_shapes( {0: node_dim + 1}, {0: edge_storage_dim}, ) + return shapes if len(sample_inputs) == len(shapes) else (*shapes, {0: node_dim}) def count_synthetic_graph_edges( @@ -1019,9 +1030,34 @@ def _supports_graph_export(model: torch.nn.Module) -> bool: return bool(model.atomic_model.supports_graph_export()) +def _spin_scheme(model_type: str | None) -> str | None: + """Return the spin scheme a model wire type implements. + + ``"native"`` treats the magnetic moment as an equivariant descriptor input + and keeps one node per atom; ``"deepspin"`` is the virtual-atom scheme. + The scheme selects the C++ backend class that serves the artifact and is + independent of the lower-forward schema it was frozen with. + + Parameters + ---------- + model_type : str or None + The serialized model wire type. + + Returns + ------- + str or None + ``"native"``, ``"deepspin"``, or ``None`` for a spin-free model. + """ + if model_type == "native_spin": + return "native" + if model_type == "spin_ener": + return "deepspin" + return None + + def _collect_metadata( model: torch.nn.Module, - is_spin: bool = False, + spin_scheme: str | None = None, lower_kind: str = "nlist", ) -> dict: """Collect metadata from the model for C++ inference. @@ -1034,7 +1070,13 @@ def _collect_metadata( The ``fitting_output_defs`` list is also included so that ``ModelOutputDef`` can be reconstructed without loading the full model. + + ``spin_scheme`` (see :func:`_spin_scheme`) is the model's spin scheme, or + ``None`` for a spin-free model; it drives both the ``is_spin`` gate on the + spin-only fields and the ``spin_scheme`` field the C++ backend factory + dispatches on. """ + is_spin = spin_scheme is not None if is_spin: fitting_output_def = model.model_output_def().def_outp else: @@ -1082,6 +1124,11 @@ def _collect_metadata( "is_spin": is_spin, } if is_spin: + # The scheme is what selects the serving backend class in C++ + # (``deepmd_create_deepspin_backend_v1``): "native" is served by + # NativeSpinPTExpt, "deepspin" by DeepSpinPTExpt. It is a property of + # the model alone, orthogonal to ``lower_input_kind`` below. + meta["spin_scheme"] = spin_scheme meta["ntypes_spin"] = model.spin.get_ntypes_spin() meta["use_spin"] = [bool(v) for v in model.spin.use_spin] # Whether multi-rank LAMMPS needs a second "with-comm" AOTI artifact @@ -1343,17 +1390,21 @@ def deserialize_to_file( ``metadata.json``. """ lower_kind = _resolve_lower_kind(model_file, data, lower_kind) - if data["model"].get("type") == "native_spin" and lower_kind != "graph": - # Native-spin models implement ONLY the NeighborGraph lower; the - # dense/nlist trace branch does not exist for them. The public freeze - # layer resolves this before calling here (see + if data["model"].get("type") == "native_spin" and lower_kind not in ( + "graph", + "dpa4c_canonical", + ): + # Native-spin models implement the NeighborGraph lower and, for an + # eligible compressed DPA4C, the compact canonical one; the dense/nlist + # trace branch does not exist for them. The public freeze layer + # resolves this before calling here (see # deepmd.pt_expt.entrypoints.main.freeze); this guard pins the # contract for direct programmatic callers with a clear error instead # of an opaque trace-time failure. raise ValueError( - "native-spin models implement only the NeighborGraph lower " - f"(got lower_kind={lower_kind!r}); use lower_kind='graph' with a " - ".pt2 output." + "native-spin models implement only the NeighborGraph and compact " + f"canonical lowers (got lower_kind={lower_kind!r}); use " + "lower_kind='graph' with a .pt2 output." ) # A graph lower deploys the fused inference pipeline. The trace runs at # DP_CUDA_INFER >= 2 so the analytic backward and CSR scatter remain custom @@ -1435,17 +1486,18 @@ def _trace_and_export( target_device = _env.DEVICE - # Detect spin model. Two flavors share the ``is_spin`` gate below (both + # Detect spin model. Two schemes share the ``is_spin`` gate below (both # need the spin-only metadata fields — ``ntypes_spin``/``use_spin`` — - # and the nlist-lower spin ABI probes), but only the NATIVE flavor + # and the nlist-lower spin ABI probes), but only the NATIVE scheme # (``native_spin``, ``NativeSpinEnergyModel``) - # rides the graph lower: the virtual-atom flavor (``spin_ener``, + # rides the graph lower: the virtual-atom scheme (``spin_ener``, # ``SpinModel``) doubles the atom count and has no graph-lower # implementation. ``is_native_spin`` distinguishes them at every seam # below (model rebuild, graph rejection, graph sample-input/dynamic-shape # ABI, trace call site). - is_native_spin = data["model"].get("type") == "native_spin" - is_spin = is_native_spin or data["model"].get("type") == "spin_ener" + spin_scheme = _spin_scheme(data["model"].get("type")) + is_native_spin = spin_scheme == "native" + is_spin = spin_scheme is not None # 1. Deserialize model on CPU for make_fx tracing. # make_fx with _allow_non_fake_inputs=True keeps real model parameters; @@ -1481,7 +1533,7 @@ def _trace_and_export( # 2. Collect metadata metadata = _collect_metadata( model, - is_spin=is_spin, + spin_scheme=spin_scheme, lower_kind=lower_kind, ) @@ -1603,6 +1655,7 @@ def _trace_and_export( model, e_sample, device=torch.device("cpu"), + want_spin=is_native_spin, ) traced = model.forward_lower_canonical_graph_exportable( *sample_inputs, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index d0f82ebe49..4ff14d03b1 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -217,17 +217,18 @@ def spin_args() -> list[Argument]: "Required for the `deepspin` scheme; ignored by the `native` scheme." ) doc_scheme = ( - "The spin implementation scheme, only effective for the DPA4/SeZM model. " + "The spin implementation scheme, only effective for descriptors that declare " + "native spin support (currently DPA4/SeZM and DPA4C). " "`native` injects the per-atom spin vector as an equivariant feature " - "(l=0 magnitude and l=1 direction) directly into the descriptor and " + "directly into the descriptor and " "derives the magnetic force as the negative spin gradient of the energy, " "without virtual atoms. `deepspin` uses the classical DeepSpin virtual-atom " "representation and is the default. Other models always use the `deepspin` scheme." ) doc_allow_missing_label = ( "Whether to admit training systems that lack a `spin` data file, filling their " - "per-atom spin with zeros instead of raising. Supported only by the SeZM/DPA4 " - "spin model; defaults to false." + "per-atom spin with zeros instead of raising. Supported only by the native " + "spin models (SeZM/DPA4 and DPA4C); defaults to false." ) return [ diff --git a/deepmd/utils/eval_metrics.py b/deepmd/utils/eval_metrics.py index d75c0b7385..622f9fad23 100644 --- a/deepmd/utils/eval_metrics.py +++ b/deepmd/utils/eval_metrics.py @@ -18,10 +18,13 @@ Callable, ) +# Full validation reports the second-rank response as stress rather than as +# virial, and stress is not a field of :class:`EnergyTypeEvalMetrics`: it needs +# the cell volume. It is therefore contributed by ``_stress_weighted_errors`` +# instead of being projected here. FULL_VALIDATION_WEIGHTED_METRIC_KEYS = { "energy_per_atom": ("mae_e_per_atom", "rmse_e_per_atom"), "force": ("mae_f", "rmse_f"), - "virial_per_atom": ("mae_v_per_atom", "rmse_v_per_atom"), } DP_TEST_WEIGHTED_METRIC_KEYS = { "energy": ("mae_e", "rmse_e"), @@ -193,6 +196,50 @@ def compute_energy_type_metrics( ) +def _stress_weighted_errors( + prediction: dict[str, np.ndarray], + test_data: dict[str, np.ndarray], + has_pbc: bool, +) -> dict[str, tuple[float, float]]: + """Return the weighted stress errors of one system. + + Stress is the negated virial divided by the cell volume, the + tensile-positive convention ``dp test`` reports. A frame whose cell is + singular carries no stress and is dropped rather than producing a + divergent entry. + + Parameters + ---------- + prediction : dict[str, np.ndarray] + Model predictions containing ``virial`` with shape ``(nframes, 9)``. + test_data : dict[str, np.ndarray] + Reference labels containing ``virial`` and ``box``, the latter with + shape ``(nframes, 9)`` in Angstrom. + has_pbc : bool + Whether the system is periodic, gating the metric. + + Returns + ------- + dict[str, tuple[float, float]] + Weighted-average-ready stress errors in eV/Angstrom^3, empty when the + system is aperiodic, carries no virial label, or has no frame with a + non-singular cell. + """ + if not (has_pbc and bool(test_data.get("find_virial", 0.0))): + return {} + box = np.asarray(test_data["box"]).reshape(-1, 3, 3) + volume = np.abs(np.linalg.det(box)) + finite = volume > 0.0 + if not np.any(finite): + return {} + scale = -1.0 / volume[finite] + stress = compute_error_stat( + prediction["virial"].reshape(-1, 9)[finite] * scale[:, None], + test_data["virial"].reshape(-1, 9)[finite] * scale[:, None], + ) + return stress.as_weighted_average_errors("mae_s", "rmse_s") + + def compute_spin_force_metrics( force_real_prediction: np.ndarray, force_real_reference: np.ndarray, @@ -290,7 +337,9 @@ def compute_full_validation_energy_metrics( Weighted-average-ready ``(value, weight)`` pairs keyed by metric. """ metrics = compute_energy_type_metrics(prediction, test_data, natoms, has_pbc) - return metrics.as_weighted_average_errors(FULL_VALIDATION_WEIGHTED_METRIC_KEYS) + errors = metrics.as_weighted_average_errors(FULL_VALIDATION_WEIGHTED_METRIC_KEYS) + errors.update(_stress_weighted_errors(prediction, test_data, has_pbc)) + return errors def compute_full_validation_spin_metrics( @@ -303,8 +352,8 @@ def compute_full_validation_spin_metrics( The energy term reuses per-atom energy errors. Forces are split into a real-atom term over all atoms and a magnetic term over the magnetic atoms - selected by ``mask_mag``. Spin models do not report virial, so no virial - metric is produced. + selected by ``mask_mag``. A periodic system additionally reports stress, + the virial divided by the cell volume. Parameters ---------- @@ -316,8 +365,7 @@ def compute_full_validation_spin_metrics( natoms : int The number of atoms per frame, used for per-atom normalization. has_pbc : bool - Unused; spin full validation never reports virial. Present to keep a - uniform profile signature. + Whether the system is periodic, gating the stress metric. Returns ------- @@ -341,6 +389,7 @@ def compute_full_validation_spin_metrics( errors.update( spin_metrics.as_weighted_average_errors(DP_TEST_SPIN_WEIGHTED_METRIC_KEYS) ) + errors.update(_stress_weighted_errors(prediction, test_data, has_pbc)) return errors @@ -411,42 +460,42 @@ class FullValidationMetricProfile: ("E_RMSE", "rmse_e_per_atom"), ("F_MAE", "mae_f"), ("F_RMSE", "rmse_f"), - ("V_MAE", "mae_v_per_atom"), - ("V_RMSE", "rmse_v_per_atom"), + ("S_MAE", "mae_s"), + ("S_RMSE", "rmse_s"), ), metric_key_map={ "e:mae": "mae_e_per_atom", "e:rmse": "rmse_e_per_atom", "f:mae": "mae_f", "f:rmse": "rmse_f", - "v:mae": "mae_v_per_atom", - "v:rmse": "rmse_v_per_atom", + "s:mae": "mae_s", + "s:rmse": "rmse_s", }, metric_family_by_key={ "mae_e_per_atom": "e", "rmse_e_per_atom": "e", "mae_f": "f", "rmse_f": "f", - "mae_v_per_atom": "v", - "rmse_v_per_atom": "v", + "mae_s": "s", + "rmse_s": "s", }, unit_by_family={ "e": ("meV/atom", 1000.0), "f": ("meV/Å", 1000.0), - "v": ("meV/atom", 1000.0), + "s": ("meV/ų", 1000.0), }, prefactor_by_metric={ "e:mae": ("start_pref_e", "limit_pref_e"), "e:rmse": ("start_pref_e", "limit_pref_e"), "f:mae": ("start_pref_f", "limit_pref_f"), "f:rmse": ("start_pref_f", "limit_pref_f"), - "v:mae": ("start_pref_v", "limit_pref_v"), - "v:rmse": ("start_pref_v", "limit_pref_v"), + "s:mae": ("start_pref_v", "limit_pref_v"), + "s:rmse": ("start_pref_v", "limit_pref_v"), }, needs_spin=False, log_header_note=( "# E uses per-atom energy, F uses component-wise force errors, " - "and V uses virial normalized by natoms.\n" + "and S uses stress, the virial divided by the cell volume.\n" ), compute_system_metrics=compute_full_validation_energy_metrics, ) @@ -460,6 +509,8 @@ class FullValidationMetricProfile: ("FR_RMSE", "rmse_fr"), ("FM_MAE", "mae_fm"), ("FM_RMSE", "rmse_fm"), + ("S_MAE", "mae_s"), + ("S_RMSE", "rmse_s"), ), metric_key_map={ "e:mae": "mae_e_per_atom", @@ -468,6 +519,8 @@ class FullValidationMetricProfile: "fr:rmse": "rmse_fr", "fm:mae": "mae_fm", "fm:rmse": "rmse_fm", + "s:mae": "mae_s", + "s:rmse": "rmse_s", }, metric_family_by_key={ "mae_e_per_atom": "e", @@ -476,11 +529,14 @@ class FullValidationMetricProfile: "rmse_fr": "fr", "mae_fm": "fm", "rmse_fm": "fm", + "mae_s": "s", + "rmse_s": "s", }, unit_by_family={ "e": ("meV/atom", 1000.0), "fr": ("meV/Å", 1000.0), "fm": ("meV/μB", 1000.0), + "s": ("meV/ų", 1000.0), }, prefactor_by_metric={ "e:mae": ("start_pref_e", "limit_pref_e"), @@ -489,11 +545,14 @@ class FullValidationMetricProfile: "fr:rmse": ("start_pref_fr", "limit_pref_fr"), "fm:mae": ("start_pref_fm", "limit_pref_fm"), "fm:rmse": ("start_pref_fm", "limit_pref_fm"), + "s:mae": ("start_pref_v", "limit_pref_v"), + "s:rmse": ("start_pref_v", "limit_pref_v"), }, needs_spin=True, log_header_note=( "# E uses per-atom energy, FR uses component-wise real-atom force " - "errors, and FM uses magnetic-atom force errors.\n" + "errors, FM uses magnetic-atom force errors, and S uses stress, the " + "virial divided by the cell volume.\n" ), compute_system_metrics=compute_full_validation_spin_metrics, ) diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md index c1a217f1af..e93faf03b8 100644 --- a/doc/model/dpa4c.md +++ b/doc/model/dpa4c.md @@ -215,6 +215,133 @@ Compression requires: `dp --pt-expt compress` reports an explicit error when the configuration falls outside these sets. +## Native spin + +DPA4C accepts a per-atom magnetic moment as an equivariant descriptor input. +The moment is not represented by a virtual atom: the atom count of the model +equals the number of physical atoms, and the magnetic force is the negative +spin gradient of the same energy that yields the conservative force, + +```math +\mathbf{F}_i = -\frac{\partial E}{\partial \mathbf{r}_i}, +\qquad +\mathbf{F}^{m}_i = -\frac{\partial E}{\partial \mathbf{s}_i} . +``` + +### Symmetry + +The magnetic moment is an axial vector: it is even under spatial inversion and +odd under time reversal, whereas a displacement is odd under inversion and even +under time reversal. The descriptor therefore emits only invariants of even +total spin order, which leaves it invariant under the full orthogonal group +acting jointly on positions and moments, including improper operations, and +invariant under time reversal. Reversing every moment leaves the energy +unchanged and reverses the magnetic force. + +Four families of spin channels are accumulated over the neighbor shell and +contracted against one another and against the geometric moments: + +| Family | Content | Interaction it represents | +| ----------------------------------- | ------------------------------------------------------------------------------------------------- | ------------------------------------------- | +| Isotropic vector | $\sum_j \varphi_c(r_{ij})\,\hat{\mathbf{s}}_j$ | Heisenberg exchange | +| Bond-projected vector | $\sum_j \varphi_c(r_{ij})\,(\hat{\mathbf{s}}_j\cdot\hat{\mathbf{u}}_{ij})\,\hat{\mathbf{u}}_{ij}$ | Symmetric anisotropic exchange | +| Quadrupole | $\sum_j \varphi_c(r_{ij})\,B_2(\hat{\mathbf{s}}_j)$ | Biquadratic exchange, single-ion anisotropy | +| Magnitude and magnetic coordination | $\sum_j \varphi_c(r_{ij})\,\lvert\mathbf{s}_j\rvert^2$ and the gated neighbor count | Longitudinal and stoichiometric terms | + +The width of the spin block follows the degree-two width of the geometric +descriptor, so it is set by {ref}`channels ` +and has no knob of its own. + +Two-body Heisenberg exchange, biquadratic exchange and single-ion anisotropy +are represented exactly rather than approximately: each corresponds to a single +emitted invariant times a learned radial profile. The Dzyaloshinskii-Moriya +interaction is not representable at any order, because the invariant read-out +contains no antisymmetric contraction. + +### Configuration + +Native spin is requested at the model level, not on the descriptor. The +`use_spin` list marks the magnetic types, either as booleans over the type map +or by element name: + +```json +{ + "model": { + "type_map": [ + "Ni", + "O" + ], + "spin": { + "scheme": "native", + "use_spin": [ + true, + false + ] + }, + "descriptor": { + "type": "dpa4c", + "rcut": 6.0, + "channels": 32, + "lmax": 2, + "precision": "float32" + } + } +} +``` + +The `native` scheme is required; the virtual-atom `deepspin` scheme is not +supported by this descriptor. Training uses the `ener_spin` loss, whose +`start_pref_fm` and `limit_pref_fm` weight the magnetic force. A complete +example is provided in `examples/spin/dpa4c/input.json`. + +A moment is conditioned by a per-type gate and a reference magnitude measured +from the training corpus, so a non-magnetic type contributes exactly zero to +every spin channel and the magnitude of a magnetic type is normalized to order +unity. A model that declares a magnetic type but receives no moment is +rejected rather than evaluated at zero, since the latter is indistinguishable +from a broken data pipeline and reports a vanishing magnetic force. + +### Running in LAMMPS + +A native-spin model is served by the `dpa4spin` pair style, and by +`dpa4spin/kk` under Kokkos. Both require `atom_style spin` and a model frozen +with the compact canonical graph lower, which compression selects on its own +for an eligible DPA4C. Freeze first and compress the frozen artifact, as for +any other DPA4C model: + +```bash +dp --pt-expt freeze -c model.ckpt.pt -o frozen_model +dp --pt-expt compress -i frozen_model.pt2 -o compressed_model.pt2 +``` + +The `lower_input_kind` of `compressed_model.pt2` reads `dpa4c_canonical`, which +is what the pair styles require; `frozen_model.pt2` alone carries the plain +graph lower and is refused. + +```lammps +atom_style spin +pair_style dpa4spin compressed_model.pt2 +pair_coeff * * Ni O +``` + +The Kokkos style keeps the graph and the moment in device memory for the whole +step. Ghost moments are supplied by the forward communication that +`atom_style spin` already performs, and the magnetic force is reduced back onto +owning atoms alongside the conservative force, so domain decomposition needs no +additional exchange: + +```bash +lmp -k on g 1 -sf kk -in in.lammps +``` + +A worked example, a rocksalt NiO cell in its type-II antiferromagnetic order, +is provided in `examples/spin/dpa4c/lmp/`. + +`min_style spin` reads the magnetic force from the pair style and relaxes the +moment directions. Spin dynamics through `fix nve/spin` requires a LAMMPS build +whose fix recognizes this pair style, because the stock fix accumulates the +magnetic force only from pair styles matching its own name pattern. + ## Export and running in LAMMPS DPA4C uses the PyTorch `.pt2` (AOTInductor) export path. Freeze with the graph @@ -332,3 +459,6 @@ positions and types. - The descriptor is one-hop local by construction. Interactions beyond `rcut` are not represented, and unlike a message-passing model the effective range cannot be extended by adding layers. +- Native spin requires the `native` scheme; the virtual-atom `deepspin` scheme + is not supported. The Dzyaloshinskii-Moriya interaction is not representable, + as explained under [Native spin](#native-spin). diff --git a/examples/spin/dpa4c/input.json b/examples/spin/dpa4c/input.json new file mode 100644 index 0000000000..0295e66f17 --- /dev/null +++ b/examples/spin/dpa4c/input.json @@ -0,0 +1,86 @@ +{ + "_comment": "DPA4C native-spin training example: the per-atom spin enters the descriptor as extra equivariant channels, and the magnetic force is the negative spin gradient of the energy. No virtual atoms are created.", + "model": { + "type_map": [ + "Ni", + "O" + ], + "spin": { + "scheme": "native", + "use_spin": [ + true, + false + ] + }, + "descriptor": { + "type": "dpa4c", + "rcut": 6.0, + "channels": 32, + "lmax": 2, + "basis_type": "bessel", + "n_radial": 16, + "radial_modes": 0, + "precision": "float32", + "seed": 42 + }, + "fitting_net": { + "neuron": [ + 192, + 192, + 192 + ], + "resnet_dt": false, + "activation_function": "silu", + "precision": "float32", + "seed": 42 + } + }, + "learning_rate": { + "type": "wsd", + "start_lr": 1e-3, + "stop_lr": 1e-6, + "warmup_ratio": 0.003, + "warmup_start_factor": 0.2, + "decay_phase_ratio": 0.2, + "decay_type": "inverse_linear" + }, + "loss": { + "type": "ener_spin", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_fr": 1000, + "limit_pref_fr": 1, + "start_pref_fm": 1000, + "limit_pref_fm": 1 + }, + "optimizer": { + "type": "HybridMuon", + "weight_decay": 0.001 + }, + "training": { + "stat_file": "./dpa4c_spin.hdf5", + "training_data": { + "systems": [ + "../data_reformat/data_0", + "../data_reformat/data_1" + ], + "batch_size": 1 + }, + "validation_data": { + "systems": [ + "../data_reformat/data_2" + ], + "batch_size": 1, + "numb_btch": 1 + }, + "numb_steps": 1000000, + "gradient_max_norm": 5.0, + "save_freq": 2000, + "max_ckpt_keep": 3, + "disp_file": "lcurve.out", + "disp_freq": 1000, + "disp_training": true, + "time_training": true, + "seed": 42 + } +} diff --git a/examples/spin/dpa4c/lmp/README.md b/examples/spin/dpa4c/lmp/README.md new file mode 100644 index 0000000000..cf7fe1ef64 --- /dev/null +++ b/examples/spin/dpa4c/lmp/README.md @@ -0,0 +1,65 @@ +# LAMMPS example for DPA4C native spin + +Runs a native-spin DPA4C model in LAMMPS through `pair_style dpa4spin`, and +through `dpa4spin/kk` under Kokkos. The magnetic moment enters the descriptor +as an equivariant input, so no virtual atoms are created and the atom count +equals the number of physical atoms. For the classical DeepSpin (virtual-atom) +scheme, see `examples/spin/lmp`; for DPA4 / SeZM native spin, see +`examples/spin/dpa4/lmp`. + +## Files + +| File | Description | +| ----------- | ---------------------------------------------------------------------------------------------------- | +| `in.lammps` | Single-point evaluation, spin relaxation and lattice dynamics of a NiO supercell. | +| `init.data` | `atom_style spin` data: rocksalt NiO, 32 magnetic Ni + 32 O, in its type-II antiferromagnetic order. | + +## Usage + +Train with the configuration in `../input.json`, compress, and freeze. The +pair style requires the compact canonical graph lower, which `--lower-kind auto` selects for a compressed DPA4C; the archive is target-specific and is not +shipped, so freeze locally: + +```bash +dp --pt-expt train ../input.json +dp --pt-expt compress -i model.ckpt.pt -o compressed.pt +dp --pt-expt freeze -c compressed.pt -o frozen_model --lower-kind auto +``` + +Run on the host: + +```bash +lmp -in in.lammps +``` + +Run device-resident under Kokkos, where the graph and the moment stay in device +memory for the whole step: + +```bash +lmp -k on g 1 -sf kk -in in.lammps +``` + +Either path runs under domain decomposition without any change. Ghost moments +arrive through the forward communication that `atom_style spin` already +performs, and the magnetic force is reduced onto owning atoms alongside the +conservative force: + +```bash +mpirun -np 2 lmp -in in.lammps +``` + +`nio.dump` holds the per-atom moment (`c_spin[1..4]`), the magnetic force +(`c_spin[5..7]`), which is non-zero on Ni and exactly zero on O, and the +conservative force (`fx fy fz`). + +The run has three stages. A single-point evaluation reports the energy and both +forces of the antiferromagnetic reference state. `min_style spin` then reads +the magnetic force from the pair style and relaxes the moment directions at +fixed positions. Finally `fix nvt` integrates positions at the relaxed magnetic +configuration, which exercises the conservative force, the neighbor rebuilds +and the domain decomposition; the thermostat does not touch the moments, so +they stay fixed through that stage. + +Spin dynamics through `fix nve/spin` requires a LAMMPS build whose fix +recognizes this pair style, because the stock fix accumulates the magnetic +force only from pair styles matching its own name pattern. diff --git a/examples/spin/dpa4c/lmp/in.lammps b/examples/spin/dpa4c/lmp/in.lammps new file mode 100644 index 0000000000..a2e39f97c5 --- /dev/null +++ b/examples/spin/dpa4c/lmp/in.lammps @@ -0,0 +1,66 @@ +# Native-spin DPA4C on rocksalt NiO in its type-II antiferromagnetic order. +# +# The magnetic moment enters the descriptor as an equivariant input and the +# magnetic force is the negative spin gradient of the energy, so no virtual +# atoms are created and the atom count equals the number of physical atoms. +# +# The pair style requires a model frozen with the compact canonical graph +# lower, which is what `dp --pt-expt freeze --lower-kind auto` selects for a +# compressed DPA4C: +# +# dp --pt-expt compress -i model.ckpt.pt -o compressed.pt +# dp --pt-expt freeze -c compressed.pt -o frozen_model --lower-kind auto +# +# Substitute `dpa4spin/kk` for the device-resident Kokkos path, which keeps the +# graph and the moment in device memory: +# +# lmp -k on g 1 -sf kk -in in.lammps + +units metal +dimension 3 +boundary p p p +atom_style spin +atom_modify map array +read_data init.data + +pair_style dpa4spin frozen_model.pt2 +pair_coeff * * Ni O + +neighbor 2.0 bin +neigh_modify every 10 delay 0 check yes + +compute mag all spin +compute spin all property/atom sp spx spy spz fmx fmy fmz +variable emag equal c_mag[5] + +thermo_style custom step temp pe v_emag press +thermo 1 +thermo_modify format float %14.8g + +dump 1 all custom 10 nio.dump id type x y z & + c_spin[1] c_spin[2] c_spin[3] c_spin[4] & + c_spin[5] c_spin[6] c_spin[7] fx fy fz +dump_modify 1 sort id format float %.12e + +# A single evaluation reports the energy, the conservative force and the +# magnetic force of the antiferromagnetic reference state. +run 0 + +# `min_style spin` reads the magnetic force from the pair style directly and +# relaxes the moment directions at fixed positions. Note that `fix nve/spin` +# accumulates the magnetic force only from pair styles whose name it matches, +# so spin dynamics needs a LAMMPS build aware of this style. +min_style spin +min_modify line spin_cubic discrete_factor 10.0 +minimize 1.0e-10 1.0e-8 500 1000 + +write_data relaxed.data + +# Lattice dynamics at the relaxed magnetic configuration. The thermostat +# integrates positions only, so the moments stay fixed while the conservative +# force, the neighbor rebuilds and the domain decomposition are exercised. +velocity all create 300.0 4928459 mom yes rot yes dist gaussian +fix 1 all nvt temp 300.0 300.0 0.1 +timestep 0.001 +thermo 10 +run 100 diff --git a/examples/spin/dpa4c/lmp/init.data b/examples/spin/dpa4c/lmp/init.data new file mode 100644 index 0000000000..db27133e33 --- /dev/null +++ b/examples/spin/dpa4c/lmp/init.data @@ -0,0 +1,80 @@ +Rocksalt NiO, type-II antiferromagnetic order, generated for the DPA4C native-spin example + +64 atoms +2 atom types + +0.0 8.34000000 xlo xhi +0.0 8.34000000 ylo yhi +0.0 8.34000000 zlo zhi + +Masses + +1 58.6934 +2 15.9994 + +Atoms # spin + +1 1 0.00000000 0.00000000 0.00000000 0.57735027 0.57735027 0.57735027 1.90000000 +2 2 2.08500000 0.00000000 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 +3 1 0.00000000 2.08500000 2.08500000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +4 2 2.08500000 2.08500000 2.08500000 0.00000000 0.00000000 1.00000000 0.00000000 +5 1 2.08500000 0.00000000 2.08500000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +6 2 4.17000000 0.00000000 2.08500000 0.00000000 0.00000000 1.00000000 0.00000000 +7 1 2.08500000 2.08500000 0.00000000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +8 2 4.17000000 2.08500000 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 +9 1 0.00000000 0.00000000 4.17000000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +10 2 2.08500000 0.00000000 4.17000000 0.00000000 0.00000000 1.00000000 0.00000000 +11 1 0.00000000 2.08500000 6.25500000 0.57735027 0.57735027 0.57735027 1.90000000 +12 2 2.08500000 2.08500000 6.25500000 0.00000000 0.00000000 1.00000000 0.00000000 +13 1 2.08500000 0.00000000 6.25500000 0.57735027 0.57735027 0.57735027 1.90000000 +14 2 4.17000000 0.00000000 6.25500000 0.00000000 0.00000000 1.00000000 0.00000000 +15 1 2.08500000 2.08500000 4.17000000 0.57735027 0.57735027 0.57735027 1.90000000 +16 2 4.17000000 2.08500000 4.17000000 0.00000000 0.00000000 1.00000000 0.00000000 +17 1 0.00000000 4.17000000 0.00000000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +18 2 2.08500000 4.17000000 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 +19 1 0.00000000 6.25500000 2.08500000 0.57735027 0.57735027 0.57735027 1.90000000 +20 2 2.08500000 6.25500000 2.08500000 0.00000000 0.00000000 1.00000000 0.00000000 +21 1 2.08500000 4.17000000 2.08500000 0.57735027 0.57735027 0.57735027 1.90000000 +22 2 4.17000000 4.17000000 2.08500000 0.00000000 0.00000000 1.00000000 0.00000000 +23 1 2.08500000 6.25500000 0.00000000 0.57735027 0.57735027 0.57735027 1.90000000 +24 2 4.17000000 6.25500000 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 +25 1 0.00000000 4.17000000 4.17000000 0.57735027 0.57735027 0.57735027 1.90000000 +26 2 2.08500000 4.17000000 4.17000000 0.00000000 0.00000000 1.00000000 0.00000000 +27 1 0.00000000 6.25500000 6.25500000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +28 2 2.08500000 6.25500000 6.25500000 0.00000000 0.00000000 1.00000000 0.00000000 +29 1 2.08500000 4.17000000 6.25500000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +30 2 4.17000000 4.17000000 6.25500000 0.00000000 0.00000000 1.00000000 0.00000000 +31 1 2.08500000 6.25500000 4.17000000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +32 2 4.17000000 6.25500000 4.17000000 0.00000000 0.00000000 1.00000000 0.00000000 +33 1 4.17000000 0.00000000 0.00000000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +34 2 6.25500000 0.00000000 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 +35 1 4.17000000 2.08500000 2.08500000 0.57735027 0.57735027 0.57735027 1.90000000 +36 2 6.25500000 2.08500000 2.08500000 0.00000000 0.00000000 1.00000000 0.00000000 +37 1 6.25500000 0.00000000 2.08500000 0.57735027 0.57735027 0.57735027 1.90000000 +38 2 8.34000000 0.00000000 2.08500000 0.00000000 0.00000000 1.00000000 0.00000000 +39 1 6.25500000 2.08500000 0.00000000 0.57735027 0.57735027 0.57735027 1.90000000 +40 2 8.34000000 2.08500000 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 +41 1 4.17000000 0.00000000 4.17000000 0.57735027 0.57735027 0.57735027 1.90000000 +42 2 6.25500000 0.00000000 4.17000000 0.00000000 0.00000000 1.00000000 0.00000000 +43 1 4.17000000 2.08500000 6.25500000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +44 2 6.25500000 2.08500000 6.25500000 0.00000000 0.00000000 1.00000000 0.00000000 +45 1 6.25500000 0.00000000 6.25500000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +46 2 8.34000000 0.00000000 6.25500000 0.00000000 0.00000000 1.00000000 0.00000000 +47 1 6.25500000 2.08500000 4.17000000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +48 2 8.34000000 2.08500000 4.17000000 0.00000000 0.00000000 1.00000000 0.00000000 +49 1 4.17000000 4.17000000 0.00000000 0.57735027 0.57735027 0.57735027 1.90000000 +50 2 6.25500000 4.17000000 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 +51 1 4.17000000 6.25500000 2.08500000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +52 2 6.25500000 6.25500000 2.08500000 0.00000000 0.00000000 1.00000000 0.00000000 +53 1 6.25500000 4.17000000 2.08500000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +54 2 8.34000000 4.17000000 2.08500000 0.00000000 0.00000000 1.00000000 0.00000000 +55 1 6.25500000 6.25500000 0.00000000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +56 2 8.34000000 6.25500000 0.00000000 0.00000000 0.00000000 1.00000000 0.00000000 +57 1 4.17000000 4.17000000 4.17000000 -0.57735027 -0.57735027 -0.57735027 1.90000000 +58 2 6.25500000 4.17000000 4.17000000 0.00000000 0.00000000 1.00000000 0.00000000 +59 1 4.17000000 6.25500000 6.25500000 0.57735027 0.57735027 0.57735027 1.90000000 +60 2 6.25500000 6.25500000 6.25500000 0.00000000 0.00000000 1.00000000 0.00000000 +61 1 6.25500000 4.17000000 6.25500000 0.57735027 0.57735027 0.57735027 1.90000000 +62 2 8.34000000 4.17000000 6.25500000 0.00000000 0.00000000 1.00000000 0.00000000 +63 1 6.25500000 6.25500000 4.17000000 0.57735027 0.57735027 0.57735027 1.90000000 +64 2 8.34000000 6.25500000 4.17000000 0.00000000 0.00000000 1.00000000 0.00000000 diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 97ecdb8dab..092cb30730 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -1255,6 +1255,59 @@ extern void DP_DeepSpinComputeNListf3(DP_DeepSpin* dp, float* atomic_energy, float* atomic_virial); +/** + * @brief Evaluate a compact canonical graph on the model device with a DP spin + * model. + * + * @param[in] dp The DP spin model to use. + * @param[out] d_atom_energy Per-local-atom energy, shape ``(nloc)``. + * @param[out] d_force Per-node force, shape ``(nall_nodes, 3)``. + * @param[out] d_force_mag Per-node magnetic force, shape ``(nall_nodes, 3)``. + * @param[out] d_atom_virial Per-node virial, shape ``(nall_nodes, 9)``. + * @param[in] d_atype Per-node atom types, shape ``(nall_nodes)``. + * @param[in] d_source Source node per edge storage slot. + * @param[in] d_edge_vec FP32 edge vectors, shape ``(edge_storage, 3)``. + * @param[in] d_destination_row_ptr Destination CSR offsets. + * @param[in] d_source_row_ptr Source CSR offsets. + * @param[in] d_source_order Source-grouped edge storage positions. + * @param[in] d_spin FP32 per-node magnetic moment, shape ``(nall_nodes, 3)``; + * halo rows carry their owner's moment. + * @param[in] nloc Number of owned local nodes. + * @param[in] nall_nodes Total local-plus-halo node count. + * @param[in] edge_storage Number of edge storage slots. + * @since API version 29 + */ +extern void DP_DeepSpinComputeCanonicalGraphGPU( + DP_DeepSpin* dp, + double* d_atom_energy, + double* d_force, + double* d_force_mag, + double* d_atom_virial, + const int64_t* d_atype, + const uint32_t* d_source, + const float* d_edge_vec, + const int64_t* d_destination_row_ptr, + const int64_t* d_source_row_ptr, + const uint32_t* d_source_order, + const float* d_spin, + int nloc, + int nall_nodes, + int64_t edge_storage); + +/** + * @brief Query whether the compact canonical graph ABI is active for a DP spin + * model. + * @since API version 29 + */ +extern bool DP_DeepSpinUsesCanonicalGraphInference(DP_DeepSpin* dp); + +/** + * @brief Query whether a DP spin model is served under the native spin scheme + * rather than the virtual-atom scheme. + * @since API version 29 + */ +extern bool DP_DeepSpinUsesNativeSpinScheme(DP_DeepSpin* dp); + /** * @brief Evaluate the energy, force and virial by using a DP with the mixed *type. (double version) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index fbff0c1823..e88a38d93a 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1889,6 +1889,51 @@ class DeepSpin : public DeepBaseModel { return dchgspin; } + /** + * @brief Evaluate a compact canonical graph on the model device. + * + * Graph, moment, and output pointers reside on the model device; the halo + * rows of the moment carry their owner's value. + */ + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_force_mag, + double* d_atom_virial, + const int64_t* d_atype, + const uint32_t* d_source, + const float* d_edge_vec, + const int64_t* d_destination_row_ptr, + const int64_t* d_source_row_ptr, + const uint32_t* d_source_order, + const float* d_spin, + const int nloc, + const int nall_nodes, + const int64_t edge_storage) { + DP_DeepSpinComputeCanonicalGraphGPU( + dp, d_atom_energy, d_force, d_force_mag, d_atom_virial, d_atype, + d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, d_spin, nloc, nall_nodes, edge_storage); + DP_CHECK_OK(DP_DeepSpinCheckOK, dp); + } + + /** + * @brief Query whether the compact canonical graph ABI is active. + */ + bool uses_canonical_graph_inference() const { + const bool result = DP_DeepSpinUsesCanonicalGraphInference(dp); + DP_CHECK_OK(DP_DeepSpinCheckOK, dp); + return result; + } + + /** + * @brief Query whether the native spin scheme is active. + */ + bool uses_native_spin_scheme() const { + const bool result = DP_DeepSpinUsesNativeSpinScheme(dp); + DP_CHECK_OK(DP_DeepSpinCheckOK, dp); + return result; + } + /** * @brief Evaluate the energy, force, magnetic force and virial by using this *DP spin model. diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index e3f9ce63f0..cfb2ed5b9e 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -2173,6 +2173,46 @@ void DP_DeepSpinComputeNListf3(DP_DeepSpin* dp, charge_spin); } +void DP_DeepSpinComputeCanonicalGraphGPU(DP_DeepSpin* dp, + double* d_atom_energy, + double* d_force, + double* d_force_mag, + double* d_atom_virial, + const int64_t* d_atype, + const uint32_t* d_source, + const float* d_edge_vec, + const int64_t* d_destination_row_ptr, + const int64_t* d_source_row_ptr, + const uint32_t* d_source_order, + const float* d_spin, + const int nloc, + const int nall_nodes, + const int64_t edge_storage) { + DP_REQUIRES_OK( + dp, dp->dp.compute_canonical_graph_gpu( + d_atom_energy, d_force, d_force_mag, d_atom_virial, d_atype, + d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, d_spin, nloc, nall_nodes, edge_storage)); +} + +bool DP_DeepSpinUsesCanonicalGraphInference(DP_DeepSpin* dp) { + try { + return dp->dp.uses_canonical_graph_inference(); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + return false; + } +} + +bool DP_DeepSpinUsesNativeSpinScheme(DP_DeepSpin* dp) { + try { + return dp->dp.uses_native_spin_scheme(); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + return false; + } +} + // end multiple frames void DP_DeepPotComputeMixedType(DP_DeepPot* dp, diff --git a/source/api_cc/CMakeLists.txt b/source/api_cc/CMakeLists.txt index fde1a99c42..fe764ab0bf 100644 --- a/source/api_cc/CMakeLists.txt +++ b/source/api_cc/CMakeLists.txt @@ -16,7 +16,8 @@ set(DEEPMD_BACKEND_IMPL_SRC ${CMAKE_CURRENT_SOURCE_DIR}/src/DeepSpinPTExpt.cc ${CMAKE_CURRENT_SOURCE_DIR}/src/DeepSpinTF.cc ${CMAKE_CURRENT_SOURCE_DIR}/src/DeepTensorPT.cc - ${CMAKE_CURRENT_SOURCE_DIR}/src/DeepTensorTF.cc) + ${CMAKE_CURRENT_SOURCE_DIR}/src/DeepTensorTF.cc + ${CMAKE_CURRENT_SOURCE_DIR}/src/NativeSpinPTExpt.cc) set(DEEPMD_BACKEND_PLUGIN_SRC ${CMAKE_CURRENT_SOURCE_DIR}/src/DeepPotJAXPlugin.cc ${CMAKE_CURRENT_SOURCE_DIR}/src/DeepPotPDPlugin.cc @@ -115,7 +116,8 @@ if(ENABLE_PYTORCH AND "${OP_CXX_ABI_PT}" EQUAL "${OP_CXX_ABI}") add_library( deepmd_backend_ptexpt SHARED - src/DeepPotPTExpt.cc src/DeepPotPTExptPlugin.cc src/DeepSpinPTExpt.cc) + src/DeepPotPTExpt.cc src/DeepPotPTExptPlugin.cc src/DeepSpinPTExpt.cc + src/NativeSpinPTExpt.cc) deepmd_configure_backend_plugin(deepmd_backend_ptexpt) target_link_libraries(deepmd_backend_ptexpt PRIVATE "${TORCH_LIBRARIES}") target_compile_definitions(deepmd_backend_ptexpt PRIVATE BUILD_PYTORCH) diff --git a/source/api_cc/include/DeepSpin.h b/source/api_cc/include/DeepSpin.h index 2a13b79cf8..14d3d685b8 100644 --- a/source/api_cc/include/DeepSpin.h +++ b/source/api_cc/include/DeepSpin.h @@ -1,6 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #pragma once +#include #include #include "DeepBaseModel.h" @@ -252,6 +253,42 @@ class DeepSpinBackend : public DeepBaseModelBackend { * Empty if the backend does not provide this information. **/ virtual std::vector get_use_spin() const { return {}; }; + + /** + * @brief GPU-resident compact canonical graph inference backend hook. + * + * Given a device-resident dual-CSR graph and the per-node moment, write the + * per-atom energy, force, magnetic force, and virial back to the device + * output pointers. The PyTorch Exportable backend overrides this; every + * other backend inherits the throwing default. The signature is torch-free + * so the dispatcher stays backend-agnostic and ``libdeepmd_cc`` need not + * link PyTorch. See DeepSpin::compute_canonical_graph_gpu for the device + * pointer and graph contracts. + */ + virtual void compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_force_mag, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const float* d_spin, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage); + virtual bool uses_canonical_graph_inference() const; + + /** + * @brief Whether the backend serves the loaded artifact under the native + * spin scheme, in which the magnetic moment is a per-node descriptor input + * and the model carries no virtual atoms. Backends of the virtual-atom + * scheme inherit the negative default. + */ + virtual bool uses_native_spin_scheme() const; }; /** @@ -544,6 +581,60 @@ class DeepSpin : public DeepBaseModel { **/ std::vector get_use_spin() const; + /** + * @brief Whether the loaded artifact uses the compact canonical graph ABI. + */ + bool uses_canonical_graph_inference() const; + + /** + * @brief Whether the loaded artifact is served under the native spin + * scheme rather than the virtual-atom scheme. + */ + bool uses_native_spin_scheme() const; + + /** + * @brief Fully device-resident inference for a compact canonical native-spin + *artifact. + * + * The native-spin twin of DeepPot::compute_canonical_graph_gpu: the same + *dual-CSR compact schema plus the per-node moment, and the magnetic force + *among the outputs. All pointers reference accelerator memory on the model + *device and every output is written device-to-device. ``edge_storage`` is + *the allocated edge capacity; the physical edge count is the last entry of + *``d_destination_row_ptr`` and the tail beyond it is ignored. + * + * @param[out] d_atom_energy Per-atom energy, [nloc]. + * @param[out] d_force Per-node force, [nall_nodes * 3] row-major. + * @param[out] d_force_mag Per-node magnetic force, [nall_nodes * 3] + *row-major. + * @param[out] d_atom_virial Per-node virial, [nall_nodes * 9] row-major. + * @param[in] d_atype Per-node atom types, [nall_nodes]. + * @param[in] d_source Source-node index per edge, [edge_storage]. + * @param[in] d_edge_vec Destination-major edge vectors, [edge_storage * 3]. + * @param[in] d_destination_row_ptr Destination CSR offsets, [nall_nodes + 1]. + * @param[in] d_source_row_ptr Source CSR offsets, [nall_nodes + 1]. + * @param[in] d_source_order Source-grouped edge positions, [edge_storage]. + * @param[in] d_spin Per-node magnetic moment, [nall_nodes * 3] row-major; + *ghost rows carry their owner's moment. + * @param[in] nloc Number of local atoms. + * @param[in] nall_nodes Graph node count (local + ghost). + * @param[in] edge_storage Allocated edge capacity. + */ + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_force_mag, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const float* d_spin, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage); + protected: std::shared_ptr dp; }; diff --git a/source/api_cc/include/NativeSpinPTExpt.h b/source/api_cc/include/NativeSpinPTExpt.h new file mode 100644 index 0000000000..6f990b24ac --- /dev/null +++ b/source/api_cc/include/NativeSpinPTExpt.h @@ -0,0 +1,370 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#ifdef BUILD_PYTORCH +// The AOTInductor package loader header is absent on some platforms (e.g. +// macOS x86_64); the native-spin compact backend is compiled out there. +#if __has_include() +#define BUILD_PT_EXPT_NATIVE_SPIN 1 +#else +#define BUILD_PT_EXPT_NATIVE_SPIN 0 +#endif + +#if BUILD_PT_EXPT_NATIVE_SPIN + +#include + +#include "DeepSpin.h" + +namespace torch::inductor { +class AOTIModelPackageLoader; +} + +namespace deepmd { + +struct GraphTensorPack; +struct CanonicalGraphTensorPack; + +/** + * @brief PyTorch Exportable (AOTInductor .pt2) backend for the native spin + * scheme. + * + * Native spin treats the magnetic moment as an equivariant descriptor input: + * there are no virtual atoms and no atom doubling, so a node of the neighbor + * graph is exactly one atom and the magnetic force is a per-node output next + * to the conservative force. The scheme, declared by the archive as + * ``spin_scheme == "native"``, is what selects this class; the lower-forward + * schema the artifact was frozen with is an internal branch of it. + * + * Two schemas are served, both owned by + * ``deepmd.pt_expt.model.native_spin_model.NativeSpinEnergyModel``: + * + * - ``lower_input_kind == "graph"``, the general NeighborGraph ABI + * (``forward_lower_graph_exportable``): the ten topology tensors, the + * per-node moment at positional index 10, then the conditional frame and + * atomic parameter tail. Any graph-lower descriptor can be frozen this way. + * - ``lower_input_kind == "dpa4c_canonical"``, the compact deployment ABI + * (``forward_lower_canonical_graph_exportable``): the eight dual-CSR graph + * tensors -- uint32 topology and float32 edge vectors -- and the moment at + * positional index 8, with no conditional tail. + * + * Three entry points share the selected forward: + * + * - the standalone host ``computew`` (builds its own neighbor list), + * - the LAMMPS host ``computew`` (consumes an ``InputNlist``), + * - :meth:`compute_canonical_graph_gpu`, which takes an already device-resident + * graph and moment and writes its outputs device-to-device. Device residency + * is what the compact ABI exists for, so that entry point requires it. + * + * A single rank folds ghost neighbours onto their local owners, so the graph + * carries ``nloc`` nodes. Domain decomposition keeps the extended + * local-plus-ghost node set instead, so ghost force and magnetic force rows + * survive to be folded onto their owners by LAMMPS reverse communication. That + * layout gives a ghost node no owner to draw intermediate features from, so a + * descriptor that exchanges them is confined to a single rank. + **/ +class NativeSpinPTExpt : public DeepSpinBackend { + public: + NativeSpinPTExpt(); + virtual ~NativeSpinPTExpt(); + NativeSpinPTExpt(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = ""); + /** + * @brief Load a native-spin .pt2 archive frozen with either graph schema. + * @param[in] model Path of the .pt2 model file. + * @param[in] gpu_rank The GPU rank. + * @param[in] file_content Unsupported for .pt2; must be empty. + **/ + void init(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = ""); + + double cutoff() const { + assert(inited); + return rcut; + }; + int numb_types() const { + assert(inited); + return ntypes; + }; + int numb_types_spin() const { + assert(inited); + return ntypes_spin; + }; + int dim_fparam() const { + assert(inited); + return dfparam; + }; + int dim_aparam() const { + assert(inited); + return daparam; + }; + // Charge/spin conditioning is rejected when a native-spin model is built, + // so no archive this backend accepts carries a non-zero width; ``init`` + // enforces it. + int dim_chg_spin() const override { return 0; }; + void get_type_map(std::string& type_map); + bool is_aparam_nall() const { return false; }; + bool has_default_fparam() const { + assert(inited); + return has_default_fparam_; + }; + std::vector get_use_spin() const override { + assert(inited); + return use_spin_; + }; + + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); + + /** + * @brief Fully device-resident inference on a compact canonical graph. + * + * All pointers reference accelerator memory on the model device and every + * output is written device-to-device. ``edge_storage`` is the allocated edge + * capacity; the physical edge count is the last entry of + * ``d_destination_row_ptr`` and the tail beyond it is ignored by the + * artifact. + * + * @param[out] d_atom_energy Per-atom energy, GPU [nloc]. + * @param[out] d_force Per-node force, GPU [nall_nodes * 3] row-major. + * @param[out] d_force_mag Per-node magnetic force, GPU [nall_nodes * 3] + * row-major. + * @param[out] d_atom_virial Per-node virial, GPU [nall_nodes * 9] row-major. + * @param[in] d_atype Per-node atom types, GPU [nall_nodes]. + * @param[in] d_source Source-node index per edge, GPU [edge_storage]. + * @param[in] d_edge_vec Destination-major edge vectors, GPU + * [edge_storage * 3]. + * @param[in] d_destination_row_ptr Destination CSR offsets, GPU + * [nall_nodes + 1]. + * @param[in] d_source_row_ptr Source CSR offsets, GPU [nall_nodes + 1]. + * @param[in] d_source_order Source-grouped edge positions, GPU + * [edge_storage]. + * @param[in] d_spin Per-node magnetic moment, GPU [nall_nodes * 3] + * row-major; ghost rows carry their owner's moment. + * @param[in] nloc Number of local atoms. + * @param[in] nall_nodes Graph node count (local + ghost). + * @param[in] edge_storage Allocated edge capacity. + */ + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_force_mag, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const float* d_spin, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage) override; + + bool uses_canonical_graph_inference() const override; + + bool uses_native_spin_scheme() const override; + + private: + /** + * @brief Evaluate with a pre-built neighbor list (LAMMPS path). + * + * The caller supplies the extended coordinates, so the cell plays no part + * here. Returns extended per-atom force and magnetic force so that LAMMPS + * reverse communication folds the ghost rows onto their owners. + **/ + template + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); + /** + * @brief Evaluate without a neighbor list: build one, then fold ghost + * contributions back onto their local owners. + **/ + template + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); + + /** + * @brief Run the nine-input compact canonical native-spin forward. + * + * Positional order: the eight dual-CSR graph tensors, then the per-node + * moment at index 8, which is the last slot -- the compact lower is traced + * without a conditional tail. + */ + std::vector run_model_canonical( + const CanonicalGraphTensorPack& graph, const torch::Tensor& spin); + + /** + * @brief Run the NeighborGraph native-spin forward. + * + * Positional order: the ten NeighborGraph tensors, the per-node moment at + * index 10, then the conditional tail -- the frame parameter and the atomic + * parameter, each present only when the model declares a non-zero width. + */ + std::vector run_model_graph(const GraphTensorPack& graph, + const torch::Tensor& spin, + const torch::Tensor& fparam, + const torch::Tensor& aparam); + + /** + * @brief Apply model-level pair exclusion, canonicalize the payload and run. + * + * The shared tail of both host paths: pair exclusion is a build-time + * transform applied exactly once here, the destination-major + * canonicalization follows, and the ABI branch decides whether the payload + * is narrowed to the compact dual-CSR form or fed to the NeighborGraph + * forward. The returned map holds the artifact's public output keys with the + * per-atom virial in its ``(N, 9)`` layout. + * + * @param[in,out] graph Graph payload for ``node_count`` nodes; consumed in + * place by the canonicalization. + * @param[in] node_count Number of graph nodes. + * @param[in] nloc Number of owned nodes, the prefix of the node axis. + * @param[in] spin Per-node moment, shape ``(node_count, 3)``. + * @param[in] fparam Frame parameter, ``dim_fparam`` values, or empty when + * the model carries a default or declares no width. + * @param[in] aparam Atomic parameter, ``nloc * dim_aparam`` values, or empty + * when the model declares no width. + */ + std::map run_graph_payload( + GraphTensorPack& graph, + const std::int64_t node_count, + const std::int64_t nloc, + const torch::Tensor& spin, + const std::vector& fparam, + const std::vector& aparam); + + /** + * @brief Bind the flat artifact outputs to their metadata key names. + */ + void extract_outputs(std::map& output_map, + const std::vector& flat_outputs); + + /** + * @brief Translate PyTorch exceptions into DeePMD-kit exceptions. + */ + void translate_error(std::function f); + + bool inited; + int ntypes; + int ntypes_spin; + int dfparam; + int daparam; + bool has_default_fparam_; + std::vector default_fparam_; + double rcut; + int gpu_id; + bool gpu_enabled; + // Which of the two schemas the loaded artifact declares: the compact + // dual-CSR ABI when true, the general NeighborGraph ABI when false. + bool canonical_abi_ = false; + // Edge-vector precision the NeighborGraph artifact was traced with, read + // from the ``graph_edge_dtype`` metadata field. The compact ABI is float32 + // by definition and validates that at load. + bool graph_edge_fp32_ = false; + // Whether the descriptor reads intermediate features of neighbouring nodes, + // which an extended-region graph cannot supply for a ghost node. + bool has_message_passing_ = false; + std::vector use_spin_; + std::vector type_map; + std::vector output_keys; // sorted internal output key names + // Device-resident (ntypes+1)^2 model-level pair-type keep table, uploaded + // once in ``init`` from the ``pair_exclude_types`` metadata field. An + // UNDEFINED tensor means no exclusion. Exclusion belongs to the graph build: + // the ingestion seam applies it exactly once and the exported lower consumes + // a pre-excluded payload. + torch::Tensor pair_exclude_table_; + // Cached LAMMPS skin topology, rebuilt whenever ``ago == 0``. The + // model-cutoff edge set is recomputed from it on-device every step. + NeighborListData nlist_data; + std::vector mapping_; + at::Tensor edge_index_tensor; // node-space edges (folded or extended) + at::Tensor edge_index_ext_tensor; // extended-atom edges, for the geometry + std::unique_ptr loader; +}; + +} // namespace deepmd + +#endif // BUILD_PT_EXPT_NATIVE_SPIN +#endif // BUILD_PYTORCH diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 388d48c3f6..1e45a24438 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -851,6 +851,25 @@ inline void remap_graph_outputs_to_dense_keys( } } +/** + * @brief Flatten the per-atom virial emitted by a compact canonical lower. + * + * The compact lower emits the per-atom virial as a three-by-three tensor, + * whereas the graph lower flattens it to nine components -- the layout the + * dense-key remap consumes. The key is absent when the artifact was traced + * without the per-atom virial, in which case the remap never reads it either. + * + * @param[in,out] output_map Output tensor map of a canonical forward. + */ +inline void flatten_canonical_atom_virial( + std::map& output_map) { + const auto entry = output_map.find("atom_virial"); + if (entry == output_map.end()) { + return; + } + entry->second = entry->second.reshape({entry->second.size(0), 9}); +} + /** * @brief Remap NeighborGraph (graph-schema) native-spin public outputs onto * the dense internal-key layout ``DeepSpinPTExpt::compute`` consumes. diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 4a85c18fe3..979c36cbfe 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -1223,6 +1223,9 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // The graph forward emits flat-N PUBLIC keys (atom_energy/energy/force/ // virial/atom_virial); rewrite them into the dense internal-key layout the // downstream extraction/fold-back expects. + if (lower_input_is_canonical_) { + deepmd::flatten_canonical_atom_virial(output_map); + } if (multi_rank) { // Extended region (N == nall_real): force is already per-extended-atom, // owned energy = sum over local atom energies, no zero-padding. Ghost @@ -1616,6 +1619,9 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // fold-back is a no-op on ghosts. // single_rank=true: the standalone (build_nlist) path is always // single-rank; there is no comm_dict / cross-rank ghost exchange here. + if (lower_input_is_canonical_) { + deepmd::flatten_canonical_atom_virial(output_map); + } deepmd::remap_graph_outputs_to_dense_keys(output_map, nloc, nall, atomic, /*single_rank=*/true); } diff --git a/source/api_cc/src/DeepPotPTExptPlugin.cc b/source/api_cc/src/DeepPotPTExptPlugin.cc index cfda4ff6f9..3499c62dcb 100644 --- a/source/api_cc/src/DeepPotPTExptPlugin.cc +++ b/source/api_cc/src/DeepPotPTExptPlugin.cc @@ -1,11 +1,62 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include "DeepPotPTExpt.h" #include "DeepSpinPTExpt.h" +#include "NativeSpinPTExpt.h" #ifdef BUILD_PYTORCH +#include +#include + #include "BackendPluginFactory.h" +#if BUILD_PT_EXPT_SPIN && BUILD_PT_EXPT_NATIVE_SPIN +#include "commonPTExpt.h" + +namespace { + +/** + * @brief Whether ``NativeSpinPTExpt`` serves an archive. + * + * The two spin schemes are served by different classes, and only the archive + * itself says which one it holds: every spin artifact declares its scheme in + * ``spin_scheme``. Anything else -- the virtual-atom scheme, an archive + * frozen before the field existed, and an archive whose metadata cannot be + * read -- stays with ``DeepSpinPTExpt``, so the diagnostic for a malformed + * file comes from the loader that owns the format rather than from this + * dispatch. + * + * A native-spin archive that declares ``has_comm_artifact`` also stays with + * ``DeepSpinPTExpt``. That flag marks the nested with-comm artifact carrying + * the per-layer ghost-feature exchange that domain decomposition needs, and + * ``DeepSpinPTExpt`` is the only class that drives it. The conjunct states a + * capability of the serving class rather than a property of any descriptor + * family, so it drops out once ``NativeSpinPTExpt`` gains a with-comm route + * of its own. + * + * The lower-forward schema takes no part here: it is an internal branch of + * whichever class the scheme selects. + */ +bool native_spin_backend_serves(const char* model) { + if (model == nullptr) { + return false; + } + try { + const auto metadata = deepmd::ptexpt::parse_json( + deepmd::ptexpt::read_zip_entry(model, "extra/metadata.json")); + const bool native_scheme = metadata.obj_val.count("spin_scheme") && + metadata["spin_scheme"].as_string() == "native"; + const bool needs_with_comm = metadata.obj_val.count("has_comm_artifact") && + metadata["has_comm_artifact"].as_bool(); + return native_scheme && !needs_with_comm; + } catch (const std::exception&) { + return false; + } +} + +} // namespace +#endif + extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, int gpu_rank, const char* file_content, @@ -34,6 +85,12 @@ extern "C" void* deepmd_create_deepspin_backend_v1( std::size_t file_content_size, char** error_message) { #if BUILD_PT_EXPT_SPIN +#if BUILD_PT_EXPT_NATIVE_SPIN + if (native_spin_backend_serves(model)) { + return deepmd::plugin::create_deepspin_backend( + model, gpu_rank, file_content, file_content_size, error_message); + } +#endif return deepmd::plugin::create_deepspin_backend( model, gpu_rank, file_content, file_content_size, error_message); #else diff --git a/source/api_cc/src/DeepSpin.cc b/source/api_cc/src/DeepSpin.cc index fdaa45af0a..28731ab0d8 100644 --- a/source/api_cc/src/DeepSpin.cc +++ b/source/api_cc/src/DeepSpin.cc @@ -465,6 +465,77 @@ std::vector DeepSpin::get_use_spin() const { return {}; } +void DeepSpinBackend::compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_force_mag, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const float* d_spin, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage) { + (void)d_atom_energy; + (void)d_force; + (void)d_force_mag; + (void)d_atom_virial; + (void)d_atype; + (void)d_source; + (void)d_edge_vec; + (void)d_destination_row_ptr; + (void)d_source_row_ptr; + (void)d_source_order; + (void)d_spin; + (void)nloc; + (void)nall_nodes; + (void)edge_storage; + throw deepmd::deepmd_exception( + "compact canonical graph inference is only supported by a compatible " + "PyTorch Exportable backend."); +} + +bool DeepSpinBackend::uses_canonical_graph_inference() const { return false; } + +bool DeepSpinBackend::uses_native_spin_scheme() const { return false; } + +void DeepSpin::compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_force_mag, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const float* d_spin, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage) { + // Backend-agnostic dispatch: backends that implement device-resident + // canonical inference override the hook, while the others inherit the + // throwing default. ``libdeepmd_cc`` does not link any backend, so the + // dispatch stays virtual rather than casting to a concrete backend type. + dp->compute_canonical_graph_gpu( + d_atom_energy, d_force, d_force_mag, d_atom_virial, d_atype, d_source, + d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, + d_spin, nloc, nall_nodes, edge_storage); +} + +bool DeepSpin::uses_canonical_graph_inference() const { + return dp->uses_canonical_graph_inference(); +} + +bool DeepSpin::uses_native_spin_scheme() const { + return dp->uses_native_spin_scheme(); +} + DeepSpinModelDevi::DeepSpinModelDevi() { inited = false; numb_models = 0; diff --git a/source/api_cc/src/NativeSpinPTExpt.cc b/source/api_cc/src/NativeSpinPTExpt.cc new file mode 100644 index 0000000000..18adbf64c1 --- /dev/null +++ b/source/api_cc/src/NativeSpinPTExpt.cc @@ -0,0 +1,1051 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#include "NativeSpinPTExpt.h" + +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT_NATIVE_SPIN +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "SimulationRegion.h" +#include "common.h" +#include "commonPT.h" +#include "commonPTExpt.h" +#include "device.h" +#include "errors.h" +#include "neighbor_list.h" + +using deepmd::ptexpt::parse_json; +using deepmd::ptexpt::read_zip_entry; + +using namespace deepmd; + +namespace { + +void synchronize_current_accelerator_stream() { +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + DPErrcheck(gpuDeviceSynchronize()); +#else + throw deepmd::deepmd_exception( + "GPU-resident inference requires a GPU-enabled DeePMD-kit build."); +#endif +} + +/** + * @brief Reject conditioning inputs the loaded artifact has no slot for. + * + * A declared width of zero means the traced lower carries no such slot, so + * silently dropping a value the caller supplied would hide a configuration + * error rather than surface it. + */ +template +void reject_unsupported_parametric_inputs(const std::vector& fparam, + const std::vector& aparam, + const int dim_fparam, + const int dim_aparam) { + if (dim_fparam == 0 && !fparam.empty()) { + throw deepmd::deepmd_exception( + "a frame parameter was supplied, but this native-spin artifact " + "declares dim_fparam=0 and its forward has no slot for one."); + } + if (dim_aparam == 0 && !aparam.empty()) { + throw deepmd::deepmd_exception( + "an atomic parameter was supplied, but this native-spin artifact " + "declares dim_aparam=0 and its forward has no slot for one."); + } +} + +/** + * @brief Build the frame-parameter input of the conditional graph tail. + * + * The artifact consumes it in double precision, shaped ``(1, dim_fparam)``. + * A model that carries a default supplies it whenever the caller passes none, + * which is how LAMMPS drives such a model. A zero width means the forward has + * no such slot, so the returned tensor is undefined and never marshalled. + */ +torch::Tensor make_fparam_tensor(const std::vector& fparam, + const std::vector& default_fparam, + const int dim_fparam, + const torch::Device& device) { + if (dim_fparam == 0) { + return torch::Tensor(); + } + const std::vector& values = fparam.empty() ? default_fparam : fparam; + if (static_cast(values.size()) != dim_fparam) { + throw deepmd::deepmd_exception( + "fparam holds " + std::to_string(values.size()) + + " values but the model expects dim_fparam=" + + std::to_string(dim_fparam) + + "; provide it explicitly or freeze the model with a default."); + } + return torch::from_blob(const_cast(values.data()), {1, dim_fparam}, + torch::TensorOptions().dtype(torch::kFloat64)) + .clone() + .to(device); +} + +/** + * @brief Build the atomic-parameter input of the conditional graph tail. + * + * The graph ABI carries the atomic parameter flat on the node axis. The + * caller supplies the owned rows; an extended-region graph has its halo rows + * zero-padded, which the owned-node mask makes inert. A zero width means the + * forward has no such slot, so the returned tensor is undefined and never + * marshalled. + */ +torch::Tensor make_aparam_tensor(const std::vector& aparam, + const int dim_aparam, + const std::int64_t node_count, + const std::int64_t nloc, + const torch::Device& device) { + if (dim_aparam == 0) { + return torch::Tensor(); + } + const auto f64_options = torch::TensorOptions().dtype(torch::kFloat64); + const at::Tensor owned = + aparam.empty() + ? torch::zeros({0}, f64_options).to(device) + : torch::from_blob(const_cast(aparam.data()), + {static_cast(aparam.size())}, + f64_options) + .clone() + .to(device); + return deepmd::extend_graph_aparam(owned, node_count, nloc, dim_aparam); +} + +} // namespace + +void NativeSpinPTExpt::translate_error(std::function f) { + try { + f(); + } catch (const c10::Error& e) { + throw deepmd::deepmd_exception( + "DeePMD-kit PyTorch Exportable backend error: " + + std::string(e.what())); + } catch (const deepmd::deepmd_exception&) { + throw; + } catch (const std::exception& e) { + throw deepmd::deepmd_exception( + "DeePMD-kit PyTorch Exportable backend error: " + + std::string(e.what())); + } +} + +NativeSpinPTExpt::NativeSpinPTExpt() : inited(false) {} + +NativeSpinPTExpt::NativeSpinPTExpt(const std::string& model, + const int& gpu_rank, + const std::string& file_content) + : inited(false) { + translate_error([&] { init(model, gpu_rank, file_content); }); +} + +NativeSpinPTExpt::~NativeSpinPTExpt() {} + +void NativeSpinPTExpt::init(const std::string& model, + const int& gpu_rank, + const std::string& file_content) { + if (inited) { + std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " + "nothing at the second call of initializer" + << std::endl; + return; + } + + // Register the deepmd_export::* schemas with torch's dispatcher before the + // AOTI module resolves the operators its compiled graph calls into. + deepmd::load_op_library(deepmd::DPBackend::PyTorchExportable); + + if (!file_content.empty()) { + throw deepmd::deepmd_exception( + "In-memory file_content loading is not supported for .pt2 models. " + "Please provide a file path instead."); + } + + const int gpu_num = torch::cuda::device_count(); + gpu_id = (gpu_num > 0) ? (gpu_rank % gpu_num) : 0; + gpu_enabled = torch::cuda::is_available(); + if (!gpu_enabled) { + std::cout << "load model from: " << model << " to cpu" << std::endl; + } else { +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + DPErrcheck(DPSetDevice(gpu_id)); +#endif + std::cout << "load model from: " << model << " to gpu " << gpu_id + << std::endl; + } + + const auto metadata = + parse_json(read_zip_entry(model, "extra/metadata.json")); + + // The lower kind selects an input schema, not a backend: both native-spin + // schemas are served here, and every other kind belongs to a different + // scheme or a different model class. + const std::string lower_input_kind = + metadata.obj_val.count("lower_input_kind") + ? metadata["lower_input_kind"].as_string() + : std::string(); + if (lower_input_kind == "graph") { + canonical_abi_ = false; + } else if (lower_input_kind == "dpa4c_canonical") { + canonical_abi_ = true; + } else if (lower_input_kind == "dpa1_canonical") { + throw deepmd::deepmd_exception( + "the dpa1_canonical compact artifact has no moment slot; freeze a " + "native-spin model as 'graph' or 'dpa4c_canonical'."); + } else { + throw deepmd::deepmd_exception( + "a native-spin artifact must declare lower_input_kind 'graph' or " + "'dpa4c_canonical', but this archive declares '" + + lower_input_kind + "'."); + } + + const int declared_fparam = metadata["dim_fparam"].as_int(); + const int declared_aparam = metadata["dim_aparam"].as_int(); + const int declared_chg_spin = metadata.obj_val.count("dim_chg_spin") + ? metadata["dim_chg_spin"].as_int() + : 0; + // Charge/spin FiLM conditioning is rejected when a native-spin model is + // built, so neither schema reserves a slot for it. + if (declared_chg_spin > 0) { + throw deepmd::deepmd_exception( + "native spin does not combine with charge/spin conditioning, but this " + "archive declares dim_chg_spin=" + + std::to_string(declared_chg_spin) + "."); + } + dfparam = declared_fparam; + daparam = declared_aparam; + has_default_fparam_ = metadata.obj_val.count("has_default_fparam") && + metadata["has_default_fparam"].as_bool(); + default_fparam_.clear(); + if (has_default_fparam_ && metadata.obj_val.count("default_fparam")) { + for (const auto& v : metadata["default_fparam"].as_array()) { + default_fparam_.push_back(v.as_double()); + } + } + + graph_edge_fp32_ = metadata.obj_val.count("graph_edge_dtype") && + metadata["graph_edge_dtype"].as_string() == "float32"; + if (canonical_abi_) { + if (!graph_edge_fp32_) { + throw deepmd::deepmd_exception( + "compact canonical graph artifacts require float32 edge vectors."); + } + if (!metadata.obj_val.count("canonical_index_dtype") || + metadata["canonical_index_dtype"].as_string() != "uint32") { + throw deepmd::deepmd_exception( + "compact canonical graph artifacts require uint32 topology; " + "re-freeze the model with the current DeePMD-kit version."); + } + // The compact lower is traced with the nine graph and moment inputs + // alone, so a model declaring any conditioning width has no slot to + // receive it. + if (dfparam > 0 || daparam > 0) { + throw deepmd::deepmd_exception( + "the compact canonical native-spin ABI has no fparam / aparam slot, " + "but this model declares dim_fparam=" + + std::to_string(dfparam) + ", dim_aparam=" + std::to_string(daparam) + + "; freeze it with the graph lower instead."); + } + } + // The per-atom virial is a structural part of both contracts: the + // device-resident entry point returns it unconditionally and the host paths + // reduce it into the global virial. + if (!metadata.obj_val.count("do_atomic_virial") || + !metadata["do_atomic_virial"].as_bool()) { + throw deepmd::deepmd_exception( + "native-spin graph artifacts must be exported with the per-atom " + "virial."); + } + has_message_passing_ = metadata.obj_val.count("has_message_passing") && + metadata["has_message_passing"].as_bool(); + + rcut = metadata["rcut"].as_double(); + ntypes = metadata.obj_val.count("ntypes") + ? metadata["ntypes"].as_int() + : static_cast(metadata["type_map"].as_array().size()); + ntypes_spin = metadata.obj_val.count("ntypes_spin") + ? metadata["ntypes_spin"].as_int() + : 0; + + use_spin_.clear(); + if (metadata.obj_val.count("use_spin")) { + for (const auto& v : metadata["use_spin"].as_array()) { + use_spin_.push_back(v.as_bool()); + } + } + + type_map.clear(); + for (const auto& v : metadata["type_map"].as_array()) { + type_map.push_back(v.as_string()); + } + + output_keys.clear(); + for (const auto& v : metadata["output_keys"].as_array()) { + output_keys.push_back(v.as_string()); + } + + loader = std::make_unique( + model, "model", false, 1, + gpu_enabled ? static_cast(gpu_id) + : static_cast(-1)); + + // Model-level pair-type exclusion keeps its table on the model device for + // the lifetime of the backend, so the per-step graph build indexes it + // without a host round trip. + { + std::vector> pair_exclude_types; + if (metadata.obj_val.count("pair_exclude_types")) { + for (const auto& v : metadata["pair_exclude_types"].as_array()) { + pair_exclude_types.emplace_back(v[0].as_int(), v[1].as_int()); + } + } + std::vector table = + deepmd::buildPairExcludeTable(ntypes, pair_exclude_types); + if (!table.empty()) { + const torch::Device device = gpu_enabled + ? torch::Device(torch::kCUDA, gpu_id) + : torch::Device(torch::kCPU); + pair_exclude_table_ = + torch::from_blob(table.data(), + {static_cast(table.size())}, + torch::TensorOptions().dtype(torch::kInt32)) + .clone() + .to(device); + } + } + + int num_intra_nthreads, num_inter_nthreads; + get_env_nthreads(num_intra_nthreads, num_inter_nthreads); + if (num_inter_nthreads) { + try { + at::set_num_interop_threads(num_inter_nthreads); + } catch (...) { + } + } + if (num_intra_nthreads) { + try { + at::set_num_threads(num_intra_nthreads); + } catch (...) { + } + } + + inited = true; +} + +void NativeSpinPTExpt::get_type_map(std::string& type_map_str) { + type_map_str.clear(); + for (const auto& t : type_map) { + if (!type_map_str.empty()) { + type_map_str += " "; + } + type_map_str += t; + } +} + +std::vector NativeSpinPTExpt::run_model_canonical( + const CanonicalGraphTensorPack& graph, const torch::Tensor& spin) { + // The moment shares the float32 precision of the compact geometry; the cast + // is a no-op for a caller that already holds a float32 tensor. + return loader->run({graph.atype, graph.n_node, graph.n_local, graph.source, + graph.edge_vec, graph.destination_row_ptr, + graph.source_row_ptr, graph.source_order, + spin.to(torch::kFloat32)}); +} + +std::vector NativeSpinPTExpt::run_model_graph( + const GraphTensorPack& graph, + const torch::Tensor& spin, + const torch::Tensor& fparam, + const torch::Tensor& aparam) { + deepmd::check_graph_aparam_flat(aparam, daparam, + "NativeSpinPTExpt::run_model_graph"); + std::vector inputs = { + graph.atype, + graph.n_node, + graph.n_local, + graph.edge_index, + graph_edge_fp32_ ? graph.edge_vec.to(torch::kFloat32) : graph.edge_vec, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + spin}; + if (dfparam > 0) { + inputs.push_back(fparam); + } + if (daparam > 0) { + inputs.push_back(aparam); + } + return loader->run(inputs); +} + +std::map NativeSpinPTExpt::run_graph_payload( + GraphTensorPack& graph, + const std::int64_t node_count, + const std::int64_t nloc, + const torch::Tensor& spin, + const std::vector& fparam, + const std::vector& aparam) { + graph.edge_mask = + deepmd::applyPairExclusion(graph.edge_index, graph.edge_mask, graph.atype, + pair_exclude_table_, ntypes); + canonicalizeGraphPayload(graph, node_count); + std::map output_map; + if (canonical_abi_) { + extract_outputs(output_map, + run_model_canonical(compactCanonicalGraph(graph), spin)); + // The compact lower reports the per-atom virial with the fitting axis + // still in place; the graph lower already drops it. + deepmd::flatten_canonical_atom_virial(output_map); + } else { + const torch::Device device = graph.atype.device(); + extract_outputs( + output_map, + run_model_graph( + graph, spin, + make_fparam_tensor(fparam, default_fparam_, dfparam, device), + make_aparam_tensor(aparam, daparam, node_count, nloc, device))); + } + return output_map; +} + +void NativeSpinPTExpt::extract_outputs( + std::map& output_map, + const std::vector& flat_outputs) { + if (flat_outputs.size() != output_keys.size()) { + throw deepmd::deepmd_exception( + "Model returned " + std::to_string(flat_outputs.size()) + + " outputs but expected " + std::to_string(output_keys.size()) + + " (from metadata.json)"); + } + for (size_t i = 0; i < output_keys.size(); ++i) { + output_map[output_keys[i]] = flat_outputs[i]; + } +} + +// ============================================================================ +// LAMMPS path: compute with a pre-built neighbor list +// ============================================================================ + +template +void NativeSpinPTExpt::compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + const torch::Device device = gpu_enabled ? torch::Device(torch::kCUDA, gpu_id) + : torch::Device(torch::kCPU); + const auto f64_options = torch::TensorOptions().dtype(torch::kFloat64); + const auto int_option = + torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt64); + const torch::ScalarType float_type = + std::is_same::value ? torch::kFloat32 : torch::kFloat64; + const int nall = static_cast(atype.size()); + const int nframes = 1; + + // Drop the atoms whose LAMMPS type maps to NULL: the model never sees them, + // and select_map scatters the results back onto the full atom list. + std::vector dcoord, dforce, dforce_mag, aparam_real, datom_energy, + datom_virial; + std::vector datype, fwd_map, bkw_map; + int nghost_real, nall_real, nloc_real; + select_real_atoms_coord(dcoord, datype, aparam_real, nghost_real, fwd_map, + bkw_map, nall_real, nloc_real, coord, atype, aparam, + nghost, ntypes, nframes, daparam, nall, + /*aparam_nall=*/false); + const int nloc = nall_real - nghost_real; + + // Domain decomposition keeps the extended local-plus-ghost node set so that + // ghost force and magnetic force rows survive for the reverse + // communication; a single rank folds ghosts onto their local owners, which + // needs the LAMMPS atom map to resolve an owner. + const bool multi_rank = (lmp_list.nprocs > 1); + if (!multi_rank && nghost > 0 && lmp_list.mapping == nullptr) { + throw deepmd::deepmd_exception( + "single-rank inference folds ghost neighbours onto their local owners " + "through the LAMMPS atom map; add 'atom_modify map yes' to the input, " + "or populate InputNlist.mapping before calling compute()."); + } + // A ghost node of the extended graph carries no owner on this rank, so a + // descriptor that reads intermediate features of neighbouring nodes has no + // source for them and would silently evaluate a truncated environment. The + // backend factory normally keeps such an archive with ``DeepSpinPTExpt``, + // which owns the with-comm route, so this guard stands for a caller that + // constructs this class directly. + if (multi_rank && has_message_passing_) { + throw deepmd::deepmd_exception( + "this native-spin artifact reads intermediate features of " + "neighbouring nodes, which domain decomposition cannot supply for a " + "ghost node; run it on a single MPI rank."); + } + + if (nall_real == 0) { + // A rank holding neither a real local atom nor a real ghost contributes + // nothing, while the exported graph requires at least one node. + ener.assign(nframes, static_cast(0)); + force.assign(static_cast(nframes) * fwd_map.size() * 3, + static_cast(0)); + force_mag.assign(static_cast(nframes) * fwd_map.size() * 3, + static_cast(0)); + virial.assign(static_cast(nframes) * 9, static_cast(0)); + if (atomic) { + atom_energy.assign(static_cast(nframes) * fwd_map.size(), + static_cast(0)); + atom_virial.assign(static_cast(nframes) * fwd_map.size() * 9, + static_cast(0)); + } + return; + } + + const std::vector coord_d(dcoord.begin(), dcoord.end()); + std::vector spin_d(static_cast(nall_real) * 3, 0.0); + for (int ii = 0; ii < nall_real; ++ii) { + for (int dd = 0; dd < 3; ++dd) { + spin_d[static_cast(ii) * 3 + dd] = + static_cast(spin[static_cast(bkw_map[ii]) * 3 + dd]); + } + } + const at::Tensor coord_Tensor = + torch::from_blob(const_cast(coord_d.data()), {nall_real, 3}, + f64_options) + .clone() + .to(device); + const at::Tensor spin_Tensor = + torch::from_blob(spin_d.data(), {nall_real, 3}, f64_options) + .clone() + .to(device); + const std::vector atype_64(datype.begin(), datype.end()); + const at::Tensor atype_Tensor = + torch::from_blob(const_cast(atype_64.data()), {nall_real}, + int_option) + .clone() + .to(device); + + // LAMMPS sets ago == 0 on every neighbor-list rebuild, so a positive ago + // means the cached skin topology is still valid. The model-cutoff edge set + // is recomputed from it on-device every step. + if (ago == 0) { + nlist_data.copy_from_nlist(lmp_list, nall - nghost); + nlist_data.shuffle_exclude_empty(fwd_map); + mapping_.resize(nall_real); + if (lmp_list.mapping) { + for (int ii = 0; ii < nall_real; ++ii) { + mapping_[ii] = fwd_map[lmp_list.mapping[bkw_map[ii]]]; + } + } else { + for (int ii = 0; ii < nall_real; ++ii) { + mapping_[ii] = ii; + } + } + const EdgeTensorPack topology = createEdgeTensors( + nlist_data.jlist, dcoord, mapping_, nloc, nall_real, device, + /*with_geometry=*/false, /*row_centers=*/&nlist_data.ilist, + /*fold_to_local=*/!multi_rank); + edge_index_tensor = topology.edge_index; + edge_index_ext_tensor = topology.edge_index_ext; + } + + const std::int64_t node_count = multi_rank ? nall_real : nloc; + const EdgeTensorPack edges = + compactEdgeTensors(edge_index_tensor, edge_index_ext_tensor, coord_Tensor, + static_cast(rcut)); + GraphTensorPack graph; + graph.atype = atype_Tensor.slice(0, 0, node_count); + graph.n_node = torch::full({1}, node_count, int_option).to(device); + graph.n_local = + torch::full({1}, static_cast(nloc), int_option).to(device); + graph.edge_index = edges.edge_index; + graph.edge_vec = edges.edge_vec; + graph.edge_mask = edges.edge_mask; + std::map output_map = run_graph_payload( + graph, node_count, nloc, spin_Tensor.slice(0, 0, node_count), + std::vector(fparam.begin(), fparam.end()), + std::vector(aparam_real.begin(), aparam_real.end())); + + // The forward emits flat per-node public keys; rewrite them into the dense + // internal-key layout the extraction below reads. The extended node set + // already carries one row per extended atom and must not be padded. + if (multi_rank) { + deepmd::remap_graph_spin_outputs_to_dense_keys_extended(output_map, nloc, + nall_real, atomic); + } else { + deepmd::remap_graph_spin_outputs_to_dense_keys(output_map, nloc, nall_real, + atomic); + } + + const torch::Tensor cpu_energy = + output_map["energy_redu"].view({-1}).to(torch::kCPU); + ener.assign(cpu_energy.data_ptr(), + cpu_energy.data_ptr() + cpu_energy.numel()); + + const torch::Tensor cpu_force = output_map["energy_derv_r"] + .squeeze(-2) + .view({-1}) + .to(float_type) + .to(torch::kCPU); + dforce.assign(cpu_force.data_ptr(), + cpu_force.data_ptr() + cpu_force.numel()); + const torch::Tensor cpu_force_mag = output_map["energy_derv_r_mag"] + .squeeze(-2) + .view({-1}) + .to(float_type) + .to(torch::kCPU); + dforce_mag.assign( + cpu_force_mag.data_ptr(), + cpu_force_mag.data_ptr() + cpu_force_mag.numel()); + const torch::Tensor cpu_virial = output_map["energy_derv_c_redu"] + .squeeze(-2) + .view({-1}) + .to(float_type) + .to(torch::kCPU); + virial.assign(cpu_virial.data_ptr(), + cpu_virial.data_ptr() + cpu_virial.numel()); + + force.resize(static_cast(nframes) * fwd_map.size() * 3); + force_mag.resize(static_cast(nframes) * fwd_map.size() * 3); + select_map(force, dforce, bkw_map, 3, nframes, fwd_map.size(), + nall_real); + select_map(force_mag, dforce_mag, bkw_map, 3, nframes, + fwd_map.size(), nall_real); + + if (atomic) { + const torch::Tensor cpu_atom_energy = + output_map["energy"].view({-1}).to(float_type).to(torch::kCPU); + datom_energy.assign( + cpu_atom_energy.data_ptr(), + cpu_atom_energy.data_ptr() + cpu_atom_energy.numel()); + const torch::Tensor cpu_atom_virial = output_map["energy_derv_c"] + .squeeze(-2) + .view({-1}) + .to(float_type) + .to(torch::kCPU); + datom_virial.assign( + cpu_atom_virial.data_ptr(), + cpu_atom_virial.data_ptr() + cpu_atom_virial.numel()); + + atom_energy.resize(static_cast(nframes) * fwd_map.size()); + atom_virial.resize(static_cast(nframes) * fwd_map.size() * 9); + select_map(atom_energy, datom_energy, bkw_map, 1, nframes, + fwd_map.size(), nall_real); + select_map(atom_virial, datom_virial, bkw_map, 9, nframes, + fwd_map.size(), nall_real); + } +} + +template void NativeSpinPTExpt::compute>( + std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); +template void NativeSpinPTExpt::compute>( + std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); + +// ============================================================================ +// Standalone path: compute without a pre-built neighbor list +// ============================================================================ + +template +void NativeSpinPTExpt::compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + const torch::Device device = gpu_enabled ? torch::Device(torch::kCUDA, gpu_id) + : torch::Device(torch::kCPU); + const auto f64_options = torch::TensorOptions().dtype(torch::kFloat64); + const torch::ScalarType float_type = + std::is_same::value ? torch::kFloat32 : torch::kFloat64; + const int nloc = static_cast(atype.size()); + const int nframes = 1; + + // === Step 1. Supply a box when the caller has none === + // An isolated cluster is embedded in an orthorhombic cell wide enough that + // no atom sees a periodic image of another. + std::vector coord_d(coord.begin(), coord.end()); + const std::vector spin_d(spin.begin(), spin.end()); + std::vector box_d(box.begin(), box.end()); + if (box_d.empty()) { + double min_x = coord_d[0], max_x = coord_d[0]; + double min_y = coord_d[1], max_y = coord_d[1]; + double min_z = coord_d[2], max_z = coord_d[2]; + for (int ii = 1; ii < nloc; ++ii) { + min_x = std::min(min_x, coord_d[ii * 3 + 0]); + max_x = std::max(max_x, coord_d[ii * 3 + 0]); + min_y = std::min(min_y, coord_d[ii * 3 + 1]); + max_y = std::max(max_y, coord_d[ii * 3 + 1]); + min_z = std::min(min_z, coord_d[ii * 3 + 2]); + max_z = std::max(max_z, coord_d[ii * 3 + 2]); + } + for (int ii = 0; ii < nloc; ++ii) { + coord_d[ii * 3 + 0] += rcut - min_x; + coord_d[ii * 3 + 1] += rcut - min_y; + coord_d[ii * 3 + 2] += rcut - min_z; + } + box_d.assign(9, 0.0); + box_d[0] = (max_x - min_x) + 2.0 * rcut; + box_d[4] = (max_y - min_y) + 2.0 * rcut; + box_d[8] = (max_z - min_z) + 2.0 * rcut; + } + + // === Step 2. Extend with ghosts and build the neighbor list === + std::vector coord_cpy_d; + std::vector atype_cpy, mapping_vec, ncell, ngcell; + { + SimulationRegion region; + region.reinitBox(&box_d[0]); + copy_coord(coord_cpy_d, atype_cpy, mapping_vec, ncell, ngcell, coord_d, + atype, static_cast(rcut), region); + } + const int nall = static_cast(coord_cpy_d.size()) / 3; + + std::vector> nlist_raw, nlist_r_cpy; + { + SimulationRegion region; + region.reinitBox(&box_d[0]); + std::vector nat_stt(3, 0), ext_stt(3), ext_end(3); + for (int dd = 0; dd < 3; ++dd) { + ext_stt[dd] = -ngcell[dd]; + ext_end[dd] = ncell[dd] + ngcell[dd]; + } + build_nlist(nlist_raw, nlist_r_cpy, coord_cpy_d, nloc, rcut, rcut, nat_stt, + ncell, ext_stt, ext_end, region, ncell); + } + + // === Step 3. Run the forward on the folded node set === + // build_nlist keys row i to center i and already cuts at rcut, so the graph + // needs no row remapping. The path is single-rank by construction: ghosts + // fold onto their local owners and the graph carries nloc nodes. + const std::vector mapping_64(mapping_vec.begin(), + mapping_vec.end()); + GraphTensorPack graph = + buildGraphTensors(nlist_raw, coord_cpy_d, atype_cpy, mapping_64, nloc, + nall, static_cast(rcut), device); + const at::Tensor spin_Tensor = + torch::from_blob(const_cast(spin_d.data()), {nloc, 3}, + f64_options) + .clone() + .to(device); + std::map output_map = + run_graph_payload(graph, nloc, nloc, spin_Tensor, + std::vector(fparam.begin(), fparam.end()), + std::vector(aparam.begin(), aparam.end())); + deepmd::remap_graph_spin_outputs_to_dense_keys(output_map, nloc, nall, + atomic); + + // === Step 4. Read the outputs and fold ghost rows onto their owners === + const torch::Tensor cpu_energy = + output_map["energy_redu"].view({-1}).to(torch::kCPU); + ener.assign(cpu_energy.data_ptr(), + cpu_energy.data_ptr() + cpu_energy.numel()); + + const torch::Tensor cpu_virial = output_map["energy_derv_c_redu"] + .squeeze(-2) + .view({-1}) + .to(float_type) + .to(torch::kCPU); + virial.assign(cpu_virial.data_ptr(), + cpu_virial.data_ptr() + cpu_virial.numel()); + + const torch::Tensor cpu_force = output_map["energy_derv_r"] + .squeeze(-2) + .view({-1}) + .to(float_type) + .to(torch::kCPU); + const std::vector extended_force( + cpu_force.data_ptr(), + cpu_force.data_ptr() + cpu_force.numel()); + fold_back(force, extended_force, mapping_vec, nloc, nall, 3, nframes); + + const torch::Tensor cpu_force_mag = output_map["energy_derv_r_mag"] + .squeeze(-2) + .view({-1}) + .to(float_type) + .to(torch::kCPU); + const std::vector extended_force_mag( + cpu_force_mag.data_ptr(), + cpu_force_mag.data_ptr() + cpu_force_mag.numel()); + fold_back(force_mag, extended_force_mag, mapping_vec, nloc, nall, 3, nframes); + + if (atomic) { + const torch::Tensor cpu_atom_energy = + output_map["energy"].view({-1}).to(float_type).to(torch::kCPU); + atom_energy.assign( + cpu_atom_energy.data_ptr(), + cpu_atom_energy.data_ptr() + cpu_atom_energy.numel()); + + const torch::Tensor cpu_atom_virial = output_map["energy_derv_c"] + .squeeze(-2) + .view({-1}) + .to(float_type) + .to(torch::kCPU); + const std::vector extended_atom_virial( + cpu_atom_virial.data_ptr(), + cpu_atom_virial.data_ptr() + cpu_atom_virial.numel()); + fold_back(atom_virial, extended_atom_virial, mapping_vec, nloc, nall, 9, + nframes); + } +} + +template void NativeSpinPTExpt::compute>( + std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); +template void NativeSpinPTExpt::compute>( + std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic); + +// ============================================================================ +// Public wrappers +// ============================================================================ + +void NativeSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam); + compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, fparam, aparam, atomic); + }); +} + +void NativeSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam); + compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, fparam, aparam, atomic); + }); +} + +void NativeSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam); + compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, nghost, inlist, ago, fparam, aparam, atomic); + }); +} + +void NativeSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + reject_unsupported_parametric_inputs(fparam, aparam, dfparam, daparam); + compute(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, nghost, inlist, ago, fparam, aparam, atomic); + }); +} + +void NativeSpinPTExpt::compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_force_mag, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const float* d_spin, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage) { + if (!canonical_abi_) { + throw deepmd::deepmd_exception( + "device-resident inference consumes the compact canonical graph; this " + "archive declares the NeighborGraph lower, which the host entry " + "points serve."); + } + if (!gpu_enabled) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu requires a CUDA device."); + } + if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2 || + static_cast(edge_storage) > + std::numeric_limits::max()) { + throw deepmd::deepmd_exception( + "invalid compact canonical graph dimensions."); + } + + translate_error([&] { + const torch::Device device(torch::kCUDA, gpu_id); + const c10::DeviceGuard device_guard(device); + const auto opt_f32 = + torch::TensorOptions().dtype(torch::kFloat32).device(device); + const auto opt_f64 = + torch::TensorOptions().dtype(torch::kFloat64).device(device); + const auto opt_i64 = + torch::TensorOptions().dtype(torch::kInt64).device(device); + const auto opt_u32 = + torch::TensorOptions().dtype(torch::kUInt32).device(device); + CanonicalGraphTensorPack graph; + graph.atype = torch::from_blob(const_cast(d_atype), + {nall_nodes}, opt_i64); + graph.n_node = torch::full({1}, nall_nodes, opt_i64); + graph.n_local = torch::full({1}, nloc, opt_i64); + graph.source = torch::from_blob(const_cast(d_source), + {edge_storage}, opt_u32); + graph.edge_vec = torch::from_blob(const_cast(d_edge_vec), + {edge_storage, 3}, opt_f32); + graph.destination_row_ptr = + torch::from_blob(const_cast(d_destination_row_ptr), + {nall_nodes + 1}, opt_i64); + graph.source_row_ptr = torch::from_blob( + const_cast(d_source_row_ptr), {nall_nodes + 1}, opt_i64); + graph.source_order = torch::from_blob( + const_cast(d_source_order), {edge_storage}, opt_u32); + const auto spin = + torch::from_blob(const_cast(d_spin), {nall_nodes, 3}, opt_f32); + + std::map output; + extract_outputs(output, run_model_canonical(graph, spin)); + auto atom_energy = output["atom_energy"] + .reshape({nall_nodes}) + .slice(0, 0, nloc) + .contiguous(); + auto force = output["force"].reshape({nall_nodes, 3}).contiguous(); + auto force_mag = output["force_mag"].reshape({nall_nodes, 3}).contiguous(); + auto atom_virial = + output["atom_virial"].reshape({nall_nodes, 9}).contiguous(); + if (nloc > 0) { + torch::from_blob(d_atom_energy, {nloc}, opt_f64).copy_(atom_energy); + } + torch::from_blob(d_force, {nall_nodes, 3}, opt_f64).copy_(force); + torch::from_blob(d_force_mag, {nall_nodes, 3}, opt_f64).copy_(force_mag); + torch::from_blob(d_atom_virial, {nall_nodes, 9}, opt_f64) + .copy_(atom_virial); + synchronize_current_accelerator_stream(); + }); +} + +bool NativeSpinPTExpt::uses_canonical_graph_inference() const { + return canonical_abi_; +} + +// The archive reaches this class only through the scheme dispatch, which +// admits nothing but ``spin_scheme == "native"``; both schemas served here +// belong to that scheme. +bool NativeSpinPTExpt::uses_native_spin_scheme() const { return true; } + +#endif diff --git a/source/lmp/compact_canonical_graph_kokkos.h b/source/lmp/compact_canonical_graph_kokkos.h new file mode 100644 index 0000000000..8584169420 --- /dev/null +++ b/source/lmp/compact_canonical_graph_kokkos.h @@ -0,0 +1,477 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// The compact canonical graph builder is available when the LAMMPS Kokkos +// package is enabled. +#ifdef LMP_KOKKOS + +#ifndef LMP_COMPACT_CANONICAL_GRAPH_KOKKOS_H +#define LMP_COMPACT_CANONICAL_GRAPH_KOKKOS_H + +#include +#include +#include +#include +#include + +#include "atom.h" +#include "atom_kokkos.h" +#include "atom_masks.h" +#include "error.h" +#include "kokkos_type.h" +#include "neigh_list_kokkos.h" +#include "neighbor.h" +#include "pointers.h" + +namespace LAMMPS_NS { + +// Device-resident model node set and compact canonical neighbor graph. +// +// A ``.pt2`` artifact frozen with the compact canonical lower consumes a +// dual-CSR topology over a contiguous range of model nodes: uint32 source +// indices, float32 edge vectors, an int64 row pointer per node in each +// direction, and a permutation that reorders the destination-major edge list +// into source-major order. This component derives all of it from the Kokkos +// device neighbor list and leaves it on the device, so a pair style hands the +// model raw device pointers and never stages the graph through the host. +// +// The node set is the model's view of the atoms. Atom types that map to no +// model element (NULL in ``pair_coeff``) are virtual: they carry no node, so +// the surviving atoms are compacted into a contiguous node index range, which +// degenerates to the identity when no type maps to NULL. A single rank folds +// every ghost onto the local atom that owns it and works on the minimum-image +// node set, which requires the box to be thicker than twice the cutoff along +// every periodic direction. Domain decomposition instead gives every real +// ghost a node of its own -- the extended node set -- and leaves the pair +// style to fold the ghost outputs onto their owners by reverse communication. +// +// build() is one function, and public, because CUDA forbids extended device +// lambdas inside non-public members: its stages cannot be split into private +// helpers. +template +class CompactCanonicalGraphKokkos : protected Pointers { + public: + typedef ArrayTypes AT; + + CompactCanonicalGraphKokkos(class LAMMPS* lmp) + : Pointers(lmp), + nloc_model(0), + nnode_model(0), + has_null_types(false), + storage_count(0), + execution_space(ExecutionSpaceFromDevice::space), + extended_nodes(false), + edge_capacity(0) {} + + // Cache the LAMMPS type (1-based) -> model type map on the device and select + // the node-set mode; ``extended`` gives every real ghost a node of its own. + // The map comes from the pair style's coeff(), so this runs once the styles + // know their model, before the first graph is built. + void setup(const std::vector& type_idx_map, bool extended) { + host_type_map = type_idx_map; + extended_nodes = extended; + has_null_types = false; + const int ntypes = static_cast(host_type_map.size()); + d_type_map = + Kokkos::View("compact_canonical:type_map", ntypes); + auto h_type_map = Kokkos::create_mirror_view(d_type_map); + for (int t = 0; t < ntypes; ++t) { + if (host_type_map[t] < 0) { + has_null_types = true; // some LAMMPS type is a virtual (NULL) atom + } + h_type_map(t) = host_type_map[t]; + } + Kokkos::deep_copy(d_type_map, h_type_map); + } + + // Rebuild the atom -> node compaction, which only moves when the neighbor + // list is rebuilt or the atom count outgrows the maps. A pair style that + // builds its own graph over the same node set calls this directly; build() + // calls it for the compact canonical graph. + void refresh_nodes() { + const int nlocal = atom->nlocal; + const int nall = atom->nlocal + atom->nghost; + + if (neighbor->ago != 0 && (int)k_loc2model.extent(0) >= nall) { + return; + } + if ((int)k_candidate_to_model.extent(0) < nall) { + k_candidate_to_model = + DAT::tdual_int_1d("compact_canonical:candidate_to_model", nall); + } + if ((int)k_loc2model.extent(0) < nall) { + k_loc2model = DAT::tdual_int_1d("compact_canonical:loc2model", nall); + k_model2loc = DAT::tdual_int_1d("compact_canonical:model2loc", nall); + } + atomKK->sync(Host, TAG_MASK | TYPE_MASK); + auto h_loc2model = k_loc2model.view_host(); + auto h_model2loc = k_model2loc.view_host(); + const int* lmp_type = atom->type; + int m = 0; + for (int i = 0; i < nlocal; ++i) { + if (host_type_map[lmp_type[i] - 1] >= 0) { + h_loc2model(i) = m; + h_model2loc(m) = i; + ++m; + } else { + h_loc2model(i) = -1; + } + } + nloc_model = m; + for (int j = nlocal; j < nall; ++j) { + if (extended_nodes && host_type_map[lmp_type[j] - 1] >= 0) { + h_loc2model(j) = m; + h_model2loc(m) = j; + ++m; + } else { + h_loc2model(j) = -1; + } + } + nnode_model = m; + + // Resolve each candidate atom to its model node once, on the host. In the + // folded representation a ghost contributes to the node of the local atom + // that owns it, so the resolution is a composition of the ownership map + // with the model map; the extended representation gives ghosts their own + // nodes and the composition degenerates to the model map. Collapsing it + // here leaves the device traversal, which visits every candidate of every + // center, with a single gather. + auto h_candidate_to_model = k_candidate_to_model.view_host(); + if (extended_nodes) { + for (int j = 0; j < nall; ++j) { + h_candidate_to_model(j) = h_loc2model(j); + } + } else { + for (int j = 0; j < nall; ++j) { + const int owner = (j < nlocal) ? j : atom->map(atom->tag[j]); + h_candidate_to_model(j) = owner < 0 ? -1 : h_loc2model(owner); + } + } + k_candidate_to_model.template modify(); + k_candidate_to_model.template sync(); + d_candidate_to_model = k_candidate_to_model.template view(); + k_loc2model.template modify(); + k_loc2model.template sync(); + d_loc2model = k_loc2model.template view(); + k_model2loc.template modify(); + k_model2loc.template sync(); + d_model2loc = k_model2loc.template view(); + } + + // Build the compact canonical graph of the current configuration from the + // Kokkos full neighbor list. Bond vectors are the center-to-neighbor + // displacements divided by ``dist_unit_cvt_factor``, the same conversion the + // pair style applies to coordinates it hands the model. + void build(class NeighList* list, + double cutoff, + double dist_unit_cvt_factor) { + refresh_nodes(); + + auto* k_list = static_cast*>(list); + const int inum = k_list->inum; + auto d_numneigh = k_list->d_numneigh; + auto d_neighbors = k_list->d_neighbors; + auto d_ilist = k_list->d_ilist; + auto loc2model = d_loc2model; + auto candidate_to_model = d_candidate_to_model; + auto model2loc = d_model2loc; + const double cutsq = cutoff * cutoff; + const double inv_dist = 1.0 / dist_unit_cvt_factor; + const int node_count_int = nnode_model; + const std::size_t node_count = static_cast(node_count_int); + const int nall = atom->nlocal + atom->nghost; + + // === Node types in the artifact's index layout === + atomKK->sync(execution_space, TYPE_MASK); + auto type = atomKK->k_type.template view(); + auto type_map = d_type_map; + if ((int)d_model_type.extent(0) < node_count_int) { + d_model_type = Kokkos::View( + "compact_canonical:model_type", nall); + } + auto model_type = d_model_type; + Kokkos::parallel_for( + "compact_canonical:node_type", + Kokkos::RangePolicy(0, node_count_int), + KOKKOS_LAMBDA(const int m) { + model_type(m) = type_map(type(model2loc(m)) - 1); + }); + + atomKK->sync(execution_space, X_MASK); + auto x = atomKK->k_x.template view(); + + if (d_destination_row_ptr.extent(0) < node_count + 1) { + d_destination_row_ptr = Kokkos::View( + "compact_canonical:destination_row_ptr", node_count + 1); + d_source_counts = Kokkos::View( + "compact_canonical:source_counts", node_count); + d_source_row_ptr = Kokkos::View( + "compact_canonical:source_row_ptr", node_count + 1); + d_source_cursor = Kokkos::View( + "compact_canonical:source_cursor", node_count); + } + Kokkos::deep_copy(d_destination_row_ptr, std::int64_t{0}); + Kokkos::deep_copy(d_source_counts, std::uint32_t{0}); + if (node_count_int == 0) { + storage_count = min_storage_edges; + return; + } + auto destination_row_ptr = d_destination_row_ptr; + auto source_counts = d_source_counts; + + // === Destination-major CSR: per-node edge count, then its prefix sum === + Kokkos::parallel_for( + "compact_canonical:count", Kokkos::RangePolicy(0, inum), + KOKKOS_LAMBDA(const int ii) { + const int i = d_ilist(ii); + const int mi = loc2model(i); + if (mi < 0) { + return; + } + const double xi = x(i, 0); + const double yi = x(i, 1); + const double zi = x(i, 2); + const int jnum = d_numneigh(i); + std::int64_t count = 0; + for (int jj = 0; jj < jnum; ++jj) { + const int j = d_neighbors(i, jj) & NEIGHMASK; + const int mj = candidate_to_model(j); + if (mj < 0) { + continue; + } + const double dx = x(j, 0) - xi; + const double dy = x(j, 1) - yi; + const double dz = x(j, 2) - zi; + if (dx * dx + dy * dy + dz * dz < cutsq) { + ++count; + } + } + destination_row_ptr(mi) = count; + }); + + Kokkos::parallel_scan( + "compact_canonical:destination_scan", + Kokkos::RangePolicy(0, node_count_int), + KOKKOS_LAMBDA(const int node, std::int64_t& update, const bool final) { + const std::int64_t count = destination_row_ptr(node); + if (final) { + destination_row_ptr(node) = update; + } + update += count; + if (final && node == node_count_int - 1) { + destination_row_ptr(node_count_int) = update; + } + }); + std::int64_t edge_count = 0; + Kokkos::deep_copy(edge_count, + Kokkos::subview(d_destination_row_ptr, node_count_int)); + storage_count = std::max(edge_count, min_storage_edges); + if (static_cast(storage_count) > + std::numeric_limits::max()) { + error->one(FLERR, + "Compact canonical graph exceeds the uint32 edge-index range"); + } + const std::size_t required = static_cast(storage_count); + if (edge_capacity < required) { + // Thermal cutoff-count fluctuations are much smaller than the historical + // 12.5% geometric-growth reserve. A 2% reserve avoids repeated allocation + // while preventing unused edge storage from retaining several GiB at + // billion-edge scale. + const std::size_t slack = required / 50 + 64; + if (required > std::numeric_limits::max() - slack) { + error->one(FLERR, "Compact canonical graph capacity overflows size_t"); + } + edge_capacity = required + slack; + d_source = Kokkos::View( + "compact_canonical:source", edge_capacity); + d_edge_vec = Kokkos::View( + "compact_canonical:edge_vec", edge_capacity * 3); + d_source_order = Kokkos::View( + "compact_canonical:source_order", edge_capacity); + } + + auto source = d_source; + auto edge_vec = d_edge_vec; + // === Destination-major fill, one warp per center === + // A thread-per-center fill writes each surviving edge at an offset private + // to its center, so the lanes of a warp scatter their twelve-byte edge + // vectors across thirty-two unrelated rows. Cooperating on one center + // instead sends consecutive survivors to consecutive slots, which coalesces + // the dominant store stream. Candidates are taken a warp at a time and each + // lane writes at its exclusive prefix within the warp, so the edge order is + // the candidate order a serial fill produces. + constexpr int neighbor_lanes = 32; // lanes cooperating on one center + using team_policy = Kokkos::TeamPolicy; + using member_type = typename team_policy::member_type; + using lane_scratch = + Kokkos::View>; + using vector_scratch = + Kokkos::View>; + const int scratch_bytes = lane_scratch::shmem_size(neighbor_lanes) + + vector_scratch::shmem_size(3 * neighbor_lanes); + Kokkos::parallel_for( + "compact_canonical:fill", + team_policy(inum, neighbor_lanes) + .set_scratch_size(0, Kokkos::PerTeam(scratch_bytes)), + KOKKOS_LAMBDA(const member_type& team) { + const int i = d_ilist(team.league_rank()); + const int mi = loc2model(i); + if (mi < 0) { + return; + } + lane_scratch node(team.team_scratch(0), neighbor_lanes); + vector_scratch vec(team.team_scratch(0), 3 * neighbor_lanes); + + const double xi = x(i, 0); + const double yi = x(i, 1); + const double zi = x(i, 2); + const int jnum = d_numneigh(i); + const int lane = team.team_rank(); + std::int64_t edge = destination_row_ptr(mi); + for (int base = 0; base < jnum; base += neighbor_lanes) { + const int jj = base + lane; + int mj = -1; + if (jj < jnum) { + const int j = d_neighbors(i, jj) & NEIGHMASK; + mj = candidate_to_model(j); + if (mj >= 0) { + const double dx = x(j, 0) - xi; + const double dy = x(j, 1) - yi; + const double dz = x(j, 2) - zi; + if (dx * dx + dy * dy + dz * dz < cutsq) { + vec(3 * lane + 0) = static_cast(dx * inv_dist); + vec(3 * lane + 1) = static_cast(dy * inv_dist); + vec(3 * lane + 2) = static_cast(dz * inv_dist); + } else { + mj = -1; + } + } + } + node(lane) = mj; + team.team_barrier(); + + // Compact the survivors of this warp of candidates: consecutive + // survivors take consecutive slots, so the stores of a warp fall in + // one contiguous span of the edge arrays. + std::int64_t kept = 0; + Kokkos::parallel_scan( + Kokkos::TeamThreadRange(team, neighbor_lanes), + [&](const int slot, std::int64_t& offset, const bool final) { + const int target = node(slot); + if (final && target >= 0) { + const std::int64_t position = edge + offset; + source(position) = static_cast(target); + edge_vec(3 * position + 0) = vec(3 * slot + 0); + edge_vec(3 * position + 1) = vec(3 * slot + 1); + edge_vec(3 * position + 2) = vec(3 * slot + 2); + Kokkos::atomic_fetch_add(&source_counts(target), + std::uint32_t{1}); + } + offset += target >= 0 ? 1 : 0; + }, + kept); + edge += kept; + team.team_barrier(); + } + }); + + // === Source-major CSR and the permutation into source order === + auto source_row_ptr = d_source_row_ptr; + Kokkos::parallel_scan( + "compact_canonical:source_scan", + Kokkos::RangePolicy(0, node_count_int), + KOKKOS_LAMBDA(const int node, std::int64_t& update, const bool final) { + const std::int64_t count = + static_cast(source_counts(node)); + if (final) { + source_row_ptr(node) = update; + } + update += count; + if (final && node == node_count_int - 1) { + source_row_ptr(node_count_int) = update; + } + }); + auto source_cursor = d_source_cursor; + Kokkos::parallel_for( + "compact_canonical:source_cursor", + Kokkos::RangePolicy(0, node_count_int), + KOKKOS_LAMBDA(const int node) { + source_cursor(node) = + static_cast(source_row_ptr(node)); + }); + auto source_order = d_source_order; + Kokkos::parallel_for( + "compact_canonical:source_scatter", + Kokkos::RangePolicy>( + 0, edge_count), + KOKKOS_LAMBDA(const std::int64_t edge) { + const auto position = Kokkos::atomic_fetch_add( + &source_cursor(source(edge)), std::uint32_t{1}); + source_order(position) = static_cast(edge); + }); + + // === Guard edges past the physical end of the graph === + if (storage_count > edge_count) { + Kokkos::parallel_for( + "compact_canonical:guards", + Kokkos::RangePolicy>( + edge_count, storage_count), + KOKKOS_LAMBDA(const std::int64_t edge) { + source(edge) = std::uint32_t{0}; + edge_vec(3 * edge + 0) = 0.0f; + edge_vec(3 * edge + 1) = 0.0f; + edge_vec(3 * edge + 2) = 0.0f; + source_order(edge) = static_cast(edge); + }); + } + } + + // === Node set, valid after refresh_nodes() === + int nloc_model; // real local model nodes; the energy is summed over these + int nnode_model; // total model nodes (== nloc_model folded; + ghost + // extended) + bool has_null_types; // some LAMMPS type maps to no model element + // LAMMPS type (1-based) -> model type, resident on the device. + Kokkos::View d_type_map; + DAT::tdual_int_1d k_model2loc; // (nall) model node index -> atom index + typename AT::t_int_1d d_loc2model; // (nall) atom -> model node index, or -1 + typename AT::t_int_1d d_model2loc; + // (nall) candidate atom -> model node index, or -1, with a ghost already + // folded onto its owner when the node set is not extended. + typename AT::t_int_1d d_candidate_to_model; + + // === Compact canonical artifact, valid after build() === + // The edge arrays are padded to ``storage_count`` rows: the traced program + // declares its edge axis with a lower bound of two, so a graph with fewer + // physical edges is completed with guard rows that carry a zero bond vector + // on node zero. + std::int64_t storage_count; + Kokkos::View d_model_type; // (nnode_model) + Kokkos::View d_source; // (storage_count) + Kokkos::View d_edge_vec; // (3 * storage_count) + Kokkos::View d_destination_row_ptr; + Kokkos::View d_source_row_ptr; + Kokkos::View d_source_order; + + private: + // Shortest edge axis the traced program accepts. + static constexpr std::int64_t min_storage_edges = 2; + + ExecutionSpace execution_space; + std::vector host_type_map; + bool extended_nodes; + std::size_t edge_capacity; + + DAT::tdual_int_1d k_loc2model; + DAT::tdual_int_1d k_candidate_to_model; + // Per-node counts are bounded by the LAMMPS neighbor limit. CSR offsets + // remain int64 so the compact graph retains its full global edge range. + Kokkos::View d_source_counts; + Kokkos::View d_source_cursor; +}; + +} // namespace LAMMPS_NS + +#endif + +#endif // LMP_KOKKOS diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index 67d604f863..e49c50843a 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -2,7 +2,6 @@ #ifdef LMP_KOKKOS #include "pair_deepmd_kokkos.h" -#include #include #include #include @@ -27,18 +26,11 @@ using namespace LAMMPS_NS; -namespace { -// Lanes cooperating on one center's candidate list in the canonical fill. -constexpr int kNeighborLanes = 32; -} // namespace - template PairDeepMDKokkos::PairDeepMDKokkos(LAMMPS* lmp) : PairDeepMD(lmp), - has_null_types(false), + compact_graph(lmp), multi_rank(false), - nloc_model(0), - nnode_model(0), edge_capacity(0), edge_vec_fp32(false), canonical_graph(false), @@ -230,123 +222,17 @@ void PairDeepMDKokkos::init_style() { comm_reverse = 3; comm_reverse_off = 9; - // Cache the LAMMPS-type -> model-type map on the device (indexed by - // type - 1); type_idx_map is populated by the base coeff(). - const int ntypes = static_cast(type_idx_map.size()); - d_type_map = Kokkos::View("deepmd/kk:type_map", ntypes); - auto h_type_map = Kokkos::create_mirror_view(d_type_map); - has_null_types = false; - for (int t = 0; t < ntypes; ++t) { - if (type_idx_map[t] < 0) { - has_null_types = true; // some LAMMPS type is a virtual (NULL) atom - } - h_type_map(t) = type_idx_map[t]; - } - Kokkos::deep_copy(d_type_map, h_type_map); -} - -template -void PairDeepMDKokkos::prepare_model_nodes() { - const int nlocal = atom->nlocal; - const int nall = atom->nlocal + atom->nghost; - - if (neighbor->ago == 0 || (int)k_loc2model.extent(0) < nall) { - if ((int)k_candidate_to_model.extent(0) < nall) { - k_candidate_to_model = - DAT::tdual_int_1d("deepmd/kk:candidate_to_model", nall); - } - if ((int)k_loc2model.extent(0) < nall) { - k_loc2model = DAT::tdual_int_1d("deepmd/kk:loc2model", nall); - k_model2loc = DAT::tdual_int_1d("deepmd/kk:model2loc", nall); - } - atomKK->sync(Host, TAG_MASK | TYPE_MASK); - auto h_loc2model = k_loc2model.view_host(); - auto h_model2loc = k_model2loc.view_host(); - const int* lmp_type = atom->type; - int m = 0; - for (int i = 0; i < nlocal; ++i) { - if (type_idx_map[lmp_type[i] - 1] >= 0) { - h_loc2model(i) = m; - h_model2loc(m) = i; - ++m; - } else { - h_loc2model(i) = -1; - } - } - nloc_model = m; - for (int j = nlocal; j < nall; ++j) { - if (multi_rank && type_idx_map[lmp_type[j] - 1] >= 0) { - h_loc2model(j) = m; - h_model2loc(m) = j; - ++m; - } else { - h_loc2model(j) = -1; - } - } - nnode_model = m; - - // Resolve each candidate atom to its model node once, on the host. In the - // folded representation a ghost contributes to the node of the local atom - // that owns it, so the resolution is a composition of the ownership map - // with the model map; the extended representation gives ghosts their own - // nodes and the composition degenerates to the model map. Collapsing it - // here leaves the device traversal, which visits every candidate of every - // center, with a single gather. - auto h_candidate_to_model = k_candidate_to_model.view_host(); - if (multi_rank) { - for (int j = 0; j < nall; ++j) { - h_candidate_to_model(j) = h_loc2model(j); - } - } else { - for (int j = 0; j < nall; ++j) { - const int owner = (j < nlocal) ? j : atom->map(atom->tag[j]); - h_candidate_to_model(j) = owner < 0 ? -1 : h_loc2model(owner); - } - } - k_candidate_to_model.template modify(); - k_candidate_to_model.template sync(); - d_candidate_to_model = k_candidate_to_model.template view(); - k_loc2model.template modify(); - k_loc2model.template sync(); - d_loc2model = k_loc2model.template view(); - k_model2loc.template modify(); - k_model2loc.template sync(); - d_model2loc = k_model2loc.template view(); - } - - atomKK->sync(execution_space, TYPE_MASK); - auto type = atomKK->k_type.template view(); - auto type_map = d_type_map; - auto model2loc = d_model2loc; - if (canonical_graph) { - if ((int)d_model_type_i64.extent(0) < nnode_model) { - d_model_type_i64 = Kokkos::View( - "deepmd/kk:model_type_i64", nall); - } - auto model_type = d_model_type_i64; - Kokkos::parallel_for( - "deepmd/kk:mtype_i64", Kokkos::RangePolicy(0, nnode_model), - KOKKOS_LAMBDA(const int m) { - model_type(m) = type_map(type(model2loc(m)) - 1); - }); - } else { - if ((int)d_model_type.extent(0) < nnode_model) { - d_model_type = - Kokkos::View("deepmd/kk:model_type", nall); - } - auto model_type = d_model_type; - Kokkos::parallel_for( - "deepmd/kk:mtype", Kokkos::RangePolicy(0, nnode_model), - KOKKOS_LAMBDA(const int m) { - model_type(m) = type_map(type(model2loc(m)) - 1); - }); - } + // The model node set is shared by both input schemas; type_idx_map is + // populated by the base coeff(). + compact_graph.setup(type_idx_map, multi_rank); } template int PairDeepMDKokkos::build_edges_device() { const int nlocal = atom->nlocal; - prepare_model_nodes(); + const int nall = atom->nlocal + atom->nghost; + compact_graph.refresh_nodes(); + const int nnode_model = compact_graph.nnode_model; // === Neighbor list and atom views on the device === NeighListKokkos* k_list = @@ -356,14 +242,27 @@ int PairDeepMDKokkos::build_edges_device() { auto d_neighbors = k_list->d_neighbors; auto d_ilist = k_list->d_ilist; - atomKK->sync(execution_space, X_MASK); + atomKK->sync(execution_space, X_MASK | TYPE_MASK); auto x = atomKK->k_x.template view(); const double cut = cutoff; const double cutsq = cut * cut; - auto loc2model = d_loc2model; - auto candidate_to_model = d_candidate_to_model; - auto model2loc = d_model2loc; + auto loc2model = compact_graph.d_loc2model; + auto candidate_to_model = compact_graph.d_candidate_to_model; + auto model2loc = compact_graph.d_model2loc; + + // === Node types in the edge-input schema's index layout === + auto type = atomKK->k_type.template view(); + auto type_map = compact_graph.d_type_map; + if ((int)d_model_type.extent(0) < nnode_model) { + d_model_type = Kokkos::View("deepmd/kk:model_type", nall); + } + auto model_type = d_model_type; + Kokkos::parallel_for( + "deepmd/kk:mtype", Kokkos::RangePolicy(0, nnode_model), + KOKKOS_LAMBDA(const int m) { + model_type(m) = type_map(type(model2loc(m)) - 1); + }); if ((int)d_edge_offset.extent(0) < nlocal + 1) { d_edge_offset = Kokkos::View( @@ -495,7 +394,7 @@ int PairDeepMDKokkos::build_edges_device() { // Compacted coordinates in model-node order for edge-input models (the graph // lower ignores coordinates); only needed when virtual atoms compact them. - if (has_null_types) { + if (compact_graph.has_null_types) { if ((int)d_coord_model.extent(0) < 3 * nnode_model) { d_coord_model = Kokkos::View("deepmd/kk:coord_model", 3 * nnode_model); @@ -513,246 +412,6 @@ int PairDeepMDKokkos::build_edges_device() { return nedge; } -template -std::int64_t PairDeepMDKokkos::build_canonical_edges_device( - CompactCanonicalGraphWorkspace& workspace) { - prepare_model_nodes(); - - auto* k_list = static_cast*>(list); - const int inum = k_list->inum; - auto d_numneigh = k_list->d_numneigh; - auto d_neighbors = k_list->d_neighbors; - auto d_ilist = k_list->d_ilist; - atomKK->sync(execution_space, X_MASK); - auto x = atomKK->k_x.template view(); - auto loc2model = d_loc2model; - auto candidate_to_model = d_candidate_to_model; - const double cutsq = cutoff * cutoff; - const double inv_dist = 1.0 / dist_unit_cvt_factor; - const int node_count_int = nnode_model; - const std::size_t node_count = static_cast(node_count_int); - - if (workspace.destination_row_ptr.extent(0) < node_count + 1) { - workspace.destination_row_ptr = Kokkos::View( - "deepmd/kk:canonical_destination_row_ptr", node_count + 1); - workspace.source_counts = Kokkos::View( - "deepmd/kk:canonical_source_counts", node_count); - workspace.source_row_ptr = Kokkos::View( - "deepmd/kk:canonical_source_row_ptr", node_count + 1); - workspace.source_cursor = Kokkos::View( - "deepmd/kk:canonical_source_cursor", node_count); - } - Kokkos::deep_copy(workspace.destination_row_ptr, std::int64_t{0}); - Kokkos::deep_copy(workspace.source_counts, std::uint32_t{0}); - if (node_count_int == 0) { - return 0; - } - auto destination_row_ptr = workspace.destination_row_ptr; - auto source_counts = workspace.source_counts; - - Kokkos::parallel_for( - "deepmd/kk:canonical_count", Kokkos::RangePolicy(0, inum), - KOKKOS_LAMBDA(const int ii) { - const int i = d_ilist(ii); - const int mi = loc2model(i); - if (mi < 0) { - return; - } - const double xi = x(i, 0); - const double yi = x(i, 1); - const double zi = x(i, 2); - const int jnum = d_numneigh(i); - std::int64_t count = 0; - for (int jj = 0; jj < jnum; ++jj) { - const int j = d_neighbors(i, jj) & NEIGHMASK; - const int mj = candidate_to_model(j); - if (mj < 0) { - continue; - } - const double dx = x(j, 0) - xi; - const double dy = x(j, 1) - yi; - const double dz = x(j, 2) - zi; - if (dx * dx + dy * dy + dz * dz < cutsq) { - ++count; - } - } - destination_row_ptr(mi) = count; - }); - - Kokkos::parallel_scan( - "deepmd/kk:canonical_destination_scan", - Kokkos::RangePolicy(0, node_count_int), - KOKKOS_LAMBDA(const int node, std::int64_t& update, const bool final) { - const std::int64_t count = destination_row_ptr(node); - if (final) { - destination_row_ptr(node) = update; - } - update += count; - if (final && node == node_count_int - 1) { - destination_row_ptr(node_count_int) = update; - } - }); - std::int64_t edge_count = 0; - Kokkos::deep_copy(edge_count, Kokkos::subview(workspace.destination_row_ptr, - node_count_int)); - const std::int64_t storage_count = std::max(edge_count, 2); - if (static_cast(storage_count) > - std::numeric_limits::max()) { - error->one(FLERR, - "Compact canonical graph exceeds the uint32 edge-index range"); - } - const std::size_t required = static_cast(storage_count); - if (workspace.edge_capacity < required) { - // Thermal cutoff-count fluctuations are much smaller than the historical - // 12.5% geometric-growth reserve. A 2% reserve avoids repeated allocation - // while preventing unused edge storage from retaining several GiB at - // billion-edge scale. - const std::size_t slack = required / 50 + 64; - if (required > std::numeric_limits::max() - slack) { - error->one(FLERR, "Compact canonical graph capacity overflows size_t"); - } - workspace.edge_capacity = required + slack; - workspace.source = Kokkos::View( - "deepmd/kk:canonical_source", workspace.edge_capacity); - workspace.edge_vec = Kokkos::View( - "deepmd/kk:canonical_edge_vec", workspace.edge_capacity * 3); - workspace.source_order = Kokkos::View( - "deepmd/kk:canonical_source_order", workspace.edge_capacity); - } - - auto source = workspace.source; - auto edge_vec = workspace.edge_vec; - // One warp per center. A thread-per-center fill writes each surviving edge - // at an offset private to its center, so the lanes of a warp scatter their - // twelve-byte edge vectors across thirty-two unrelated rows. Cooperating on - // one center instead sends consecutive survivors to consecutive slots, which - // coalesces the dominant store stream. Candidates are taken a warp at a - // time and each lane writes at its exclusive prefix within the warp, so the - // edge order is the candidate order the serial fill produces. - using team_policy = Kokkos::TeamPolicy; - using member_type = typename team_policy::member_type; - using lane_scratch = - Kokkos::View>; - using vector_scratch = - Kokkos::View>; - const int scratch_bytes = lane_scratch::shmem_size(kNeighborLanes) + - vector_scratch::shmem_size(3 * kNeighborLanes); - Kokkos::parallel_for( - "deepmd/kk:canonical_fill", - team_policy(inum, kNeighborLanes) - .set_scratch_size(0, Kokkos::PerTeam(scratch_bytes)), - KOKKOS_LAMBDA(const member_type& team) { - const int i = d_ilist(team.league_rank()); - const int mi = loc2model(i); - if (mi < 0) { - return; - } - lane_scratch node(team.team_scratch(0), kNeighborLanes); - vector_scratch vec(team.team_scratch(0), 3 * kNeighborLanes); - - const double xi = x(i, 0); - const double yi = x(i, 1); - const double zi = x(i, 2); - const int jnum = d_numneigh(i); - const int lane = team.team_rank(); - std::int64_t edge = destination_row_ptr(mi); - for (int base = 0; base < jnum; base += kNeighborLanes) { - const int jj = base + lane; - int mj = -1; - if (jj < jnum) { - const int j = d_neighbors(i, jj) & NEIGHMASK; - mj = candidate_to_model(j); - if (mj >= 0) { - const double dx = x(j, 0) - xi; - const double dy = x(j, 1) - yi; - const double dz = x(j, 2) - zi; - if (dx * dx + dy * dy + dz * dz < cutsq) { - vec(3 * lane + 0) = static_cast(dx * inv_dist); - vec(3 * lane + 1) = static_cast(dy * inv_dist); - vec(3 * lane + 2) = static_cast(dz * inv_dist); - } else { - mj = -1; - } - } - } - node(lane) = mj; - team.team_barrier(); - - // Compact the survivors of this warp of candidates: consecutive - // survivors take consecutive slots, so the stores of a warp fall in - // one contiguous span of the edge arrays. - std::int64_t kept = 0; - Kokkos::parallel_scan( - Kokkos::TeamThreadRange(team, kNeighborLanes), - [&](const int slot, std::int64_t& offset, const bool final) { - const int target = node(slot); - if (final && target >= 0) { - const std::int64_t position = edge + offset; - source(position) = static_cast(target); - edge_vec(3 * position + 0) = vec(3 * slot + 0); - edge_vec(3 * position + 1) = vec(3 * slot + 1); - edge_vec(3 * position + 2) = vec(3 * slot + 2); - Kokkos::atomic_fetch_add(&source_counts(target), - std::uint32_t{1}); - } - offset += target >= 0 ? 1 : 0; - }, - kept); - edge += kept; - team.team_barrier(); - } - }); - - auto source_row_ptr = workspace.source_row_ptr; - Kokkos::parallel_scan( - "deepmd/kk:canonical_source_scan", - Kokkos::RangePolicy(0, node_count_int), - KOKKOS_LAMBDA(const int node, std::int64_t& update, const bool final) { - const std::int64_t count = - static_cast(source_counts(node)); - if (final) { - source_row_ptr(node) = update; - } - update += count; - if (final && node == node_count_int - 1) { - source_row_ptr(node_count_int) = update; - } - }); - auto source_cursor = workspace.source_cursor; - Kokkos::parallel_for( - "deepmd/kk:canonical_source_cursor", - Kokkos::RangePolicy(0, node_count_int), - KOKKOS_LAMBDA(const int node) { - source_cursor(node) = static_cast(source_row_ptr(node)); - }); - auto source_order = workspace.source_order; - Kokkos::parallel_for( - "deepmd/kk:canonical_source_scatter", - Kokkos::RangePolicy>( - 0, edge_count), - KOKKOS_LAMBDA(const std::int64_t edge) { - const auto position = Kokkos::atomic_fetch_add( - &source_cursor(source(edge)), std::uint32_t{1}); - source_order(position) = static_cast(edge); - }); - if (storage_count > edge_count) { - Kokkos::parallel_for( - "deepmd/kk:canonical_guards", - Kokkos::RangePolicy>( - edge_count, storage_count), - KOKKOS_LAMBDA(const std::int64_t edge) { - source(edge) = std::uint32_t{0}; - edge_vec(3 * edge + 0) = 0.0f; - edge_vec(3 * edge + 1) = 0.0f; - edge_vec(3 * edge + 2) = 0.0f; - source_order(edge) = static_cast(edge); - }); - } - return edge_count; -} - template void PairDeepMDKokkos::compute(int eflag, int vflag) { if (!device_path_ok) { @@ -776,14 +435,14 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { memoryKK->create_kokkos(k_eatom, eatom, maxeatom, "deepmd/kk:eatom"); d_eatom = k_eatom.template view(); } - std::int64_t nedge = 0; + int nedge = 0; if (canonical_graph) { - nedge = build_canonical_edges_device(canonical_workspace); + compact_graph.build(list, cutoff, dist_unit_cvt_factor); } else { nedge = build_edges_device(); } - const int nloc_m = nloc_model; // local model nodes (energy) - const int nnode_m = nnode_model; // total model nodes (force / virial) + const int nloc_m = compact_graph.nloc_model; // local nodes (energy) + const int nnode_m = compact_graph.nnode_model; // all nodes (force / virial) // Energy is per local node; force / virial span the model node set, which is // the local atoms (folded) or local + real ghost atoms (extended, up to @@ -839,8 +498,8 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // ``aparam`` is built in LAMMPS local order; when virtual atoms drop nodes it // must be compacted into model-node order (the first ``nloc_model`` nodes) so // it aligns with the atoms the model consumes. - if (has_null_types && dim_aparam > 0 && !aparam_step.empty()) { - auto h_m2l = k_model2loc.view_host(); + if (compact_graph.has_null_types && dim_aparam > 0 && !aparam_step.empty()) { + auto h_m2l = compact_graph.k_model2loc.view_host(); std::vector aparam_model(static_cast(nloc_m) * dim_aparam); for (int m = 0; m < nloc_m; ++m) { @@ -860,7 +519,7 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // compact the node set; otherwise the extended edge-input path is rejected. deepmd_compat::InputNlist comm_list; const deepmd_compat::InputNlist* comm_ptr = nullptr; - if (multi_rank && !has_null_types) { + if (multi_rank && !compact_graph.has_null_types) { comm_list = make_comm_nlist(); comm_ptr = &comm_list; } @@ -876,29 +535,30 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // Coordinates are model-node order: the compacted buffer when virtual atoms // are present, else the local coordinates directly (the graph lower ignores // them; edge-input models consume them). - const double* coord_ptr = has_null_types ? d_coord_model.data() : x.data(); + const double* coord_ptr = + compact_graph.has_null_types ? d_coord_model.data() : x.data(); try { if (canonical_graph) { - const std::int64_t storage_count = std::max(nedge, 2); - auto& workspace = canonical_workspace; deep_pot.compute_canonical_graph_gpu( d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), - d_model_type_i64.data(), workspace.source.data(), - workspace.edge_vec.data(), workspace.destination_row_ptr.data(), - workspace.source_row_ptr.data(), workspace.source_order.data(), - nloc_m, nnode_m, storage_count); + compact_graph.d_model_type.data(), compact_graph.d_source.data(), + compact_graph.d_edge_vec.data(), + compact_graph.d_destination_row_ptr.data(), + compact_graph.d_source_row_ptr.data(), + compact_graph.d_source_order.data(), nloc_m, nnode_m, + compact_graph.storage_count); } else if (edge_vec_fp32) { - deep_pot.compute_edges_gpu( - d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), - coord_ptr, d_model_type.data(), d_edge_index.data(), - d_edge_vec_float.data(), nloc_m, static_cast(nedge), fparam, - aparam_step, nnode_m, comm_ptr); + deep_pot.compute_edges_gpu(d_atom_energy.data(), d_out_force.data(), + d_atom_virial.data(), coord_ptr, + d_model_type.data(), d_edge_index.data(), + d_edge_vec_float.data(), nloc_m, nedge, + fparam, aparam_step, nnode_m, comm_ptr); } else { - deep_pot.compute_edges_gpu( - d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), - coord_ptr, d_model_type.data(), d_edge_index.data(), - d_edge_vec.data(), nloc_m, static_cast(nedge), fparam, - aparam_step, nnode_m, comm_ptr); + deep_pot.compute_edges_gpu(d_atom_energy.data(), d_out_force.data(), + d_atom_virial.data(), coord_ptr, + d_model_type.data(), d_edge_index.data(), + d_edge_vec.data(), nloc_m, nedge, fparam, + aparam_step, nnode_m, comm_ptr); } } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); @@ -915,7 +575,7 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // The scatter remains device-resident. If LAMMPS selects classic host // communication, the host pack/unpack methods synchronize the force DualView // and the completed fold is copied back once after all communication stages. - auto model2loc = d_model2loc; + auto model2loc = compact_graph.d_model2loc; const double fscale = scale[1][1] * force_unit_cvt_factor; reverse_virial = false; reverse_used_host = false; @@ -1003,7 +663,7 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // because the KOKKOS full-list path runs with newton pair disabled. auto h_av = Kokkos::create_mirror_view(d_atom_virial); Kokkos::deep_copy(h_av, d_atom_virial); - auto h_m2l = k_model2loc.view_host(); + auto h_m2l = compact_graph.k_model2loc.view_host(); const double vscale = scale[1][1] * ener_unit_cvt_factor; const int map9[9] = {0, 4, 8, 3, 6, 7, 1, 2, 5}; for (int m = 0; m < nloc_m; ++m) { diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index 93cfe4c358..4a620d19d9 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -20,6 +20,7 @@ PairStyle(deepmd/kk/host,PairDeepMDKokkos); #include #include +#include "compact_canonical_graph_kokkos.h" #include "kokkos_base.h" #include "kokkos_type.h" #include "neigh_list_kokkos.h" @@ -27,20 +28,6 @@ PairStyle(deepmd/kk/host,PairDeepMDKokkos); namespace LAMMPS_NS { -template -struct CompactCanonicalGraphWorkspace { - Kokkos::View source; - Kokkos::View edge_vec; - Kokkos::View destination_row_ptr; - // Per-node counts are bounded by the LAMMPS neighbor limit. CSR offsets - // remain int64 so the compact graph retains its full global edge range. - Kokkos::View source_counts; - Kokkos::View source_row_ptr; - Kokkos::View source_cursor; - Kokkos::View source_order; - std::size_t edge_capacity = 0; -}; - // GPU-resident inference for exported ``.pt2`` models whose forward consumes // an explicit edge graph: both the graph-input form (a compact, unpadded // neighbor graph) and the edge-input form. Both are dispatched through @@ -82,41 +69,17 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { DAT::tdual_int_1d, DAT::tdual_double_1d&) override; - // Build the device edge graph from the Kokkos full neighbor list, returning - // the edge count. A single rank folds ghosts onto local owners (minimum - // image); domain decomposition keeps the extended local-plus-ghost node set. - // Public because it launches extended device lambdas, which CUDA forbids - // inside non-public members. - void prepare_model_nodes(); + // Build the device edge graph of the edge-input schema from the Kokkos full + // neighbor list, returning the edge count. Public because it launches + // extended device lambdas, which CUDA forbids inside non-public members. int build_edges_device(); - std::int64_t build_canonical_edges_device( - CompactCanonicalGraphWorkspace& workspace); protected: - // LAMMPS type (1-based) -> model type, resident on the device. - Kokkos::View d_type_map; + // Model node set and, for a compact canonical artifact, the graph itself. + CompactCanonicalGraphKokkos compact_graph; Kokkos::View - d_model_type; // (nnode_model) type per model node - Kokkos::View - d_model_type_i64; // compact canonical artifact type per model node - // Virtual-atom (NULL type) compaction, rebuilt with the neighbor list: the - // model sees only the local atoms with a real model type, so ``model2loc`` - // lists those local indices and ``loc2model`` inverts it (-1 for virtual). - // When no type maps to NULL the compaction is the identity. - bool has_null_types; - bool multi_rank; // domain-decomposed run -> extended (local+ghost) node set - int nloc_model; // real local model nodes; the energy is summed over these - int nnode_model; // total model nodes (== nloc_model folded; + ghost - // extended) - // (nall) candidate atom -> model node index, or -1. Folding a ghost onto its - // owner and mapping that atom to a node is resolved once per neighbor - // rebuild so that the graph traversal needs a single gather per candidate. - DAT::tdual_int_1d k_candidate_to_model; - typename AT::t_int_1d d_candidate_to_model; - DAT::tdual_int_1d k_loc2model; // (nall) atom -> model node index, or -1 - DAT::tdual_int_1d k_model2loc; // (nall) model node index -> atom index - typename AT::t_int_1d d_loc2model; - typename AT::t_int_1d d_model2loc; + d_model_type; // (nnode_model) edge-input type per model node + bool multi_rank; // domain-decomposed run -> extended (local+ghost) node set Kokkos::View d_coord_model; // (3 * nnode_model), NULL case @@ -127,7 +90,6 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { Kokkos::View d_edge_vec; // (3 * nedge) Kokkos::View d_edge_vec_float; // (3 * nedge), compressed graph ABI - CompactCanonicalGraphWorkspace canonical_workspace; // Model outputs on the device. Energy is per local atom; force and virial // span the model node set (up to ``nall`` under domain decomposition). diff --git a/source/lmp/pair_dpa4spin.cpp b/source/lmp/pair_dpa4spin.cpp new file mode 100644 index 0000000000..b68b7622f3 --- /dev/null +++ b/source/lmp/pair_dpa4spin.cpp @@ -0,0 +1,517 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#include "pair_dpa4spin.h" + +#include +#include +#include +#include +#include +#include + +#include "atom.h" +#include "citeme.h" +#include "comm.h" +#include "deepmd_version.h" +#include "domain.h" +#include "error.h" +#include "force.h" +#include "memory.h" +#include "neigh_list.h" +#include "neigh_request.h" +#include "neighbor.h" +#include "update.h" +#include "utils.h" + +using namespace LAMMPS_NS; + +static const char cite_user_deepmd_package[] = + "USER-DEEPMD package:\n\n" + "@article{Wang_ComputPhysCommun_2018_v228_p178,\n" + " author = {Wang, Han and Zhang, Linfeng and Han, Jiequn and E, Weinan},\n" + " doi = {10.1016/j.cpc.2018.03.016},\n" + " url = {https://doi.org/10.1016/j.cpc.2018.03.016},\n" + " year = 2018,\n" + " month = {jul},\n" + " publisher = {Elsevier {BV}},\n" + " volume = 228,\n" + " journal = {Comput. Phys. Commun.},\n" + " title = {{DeePMD-kit: A deep learning package for many-body potential " + "energy representation and molecular dynamics}},\n" + " pages = {178--184}\n" + "}\n" + "@article{Zeng_JChemPhys_2023_v159_p054801,\n" + " title = {{DeePMD-kit v2: A software package for deep potential " + "models}},\n" + " author = {Jinzhe Zeng and Duo Zhang and Denghui Lu and Pinghui Mo and " + "Zeyu Li\n" + " and Yixiao Chen and Mari{\\'a}n Rynik and Li'ang Huang and Ziyao " + "Li and \n" + " Shaochen Shi and Yingze Wang and Haotian Ye and Ping Tuo and " + "Jiabin\n" + " Yang and Ye Ding and Yifan Li and Davide Tisi and Qiyu Zeng and " + "Han \n" + " Bao and Yu Xia and Jiameng Huang and Koki Muraoka and Yibo Wang " + "and \n" + " Junhan Chang and Fengbo Yuan and Sigbj{\\o}rn L{\\o}land Bore " + "and " + "Chun\n" + " Cai and Yinnian Lin and Bo Wang and Jiayan Xu and Jia-Xin Zhu " + "and \n" + " Chenxing Luo and Yuzhi Zhang and Rhys E A Goodall and Wenshuo " + "Liang\n" + " and Anurag Kumar Singh and Sikai Yao and Jingchao Zhang and " + "Renata\n" + " Wentzcovitch and Jiequn Han and Jie Liu and Weile Jia and Darrin " + "M\n" + " York and Weinan E and Roberto Car and Linfeng Zhang and Han " + "Wang},\n" + " journal = {J. Chem. Phys.},\n" + " volume = 159,\n" + " issue = 5, \n" + " year = 2023,\n" + " pages = 054801,\n" + " doi = {10.1063/5.0155600},\n" + "}\n" + "@Article{Zeng_JChemTheoryComput_2025_v21_p4375,\n" + " author = {Jinzhe Zeng and Duo Zhang and Anyang Peng and Xiangyu " + "Zhang and Sensen\n" + " He and Yan Wang and Xinzijian Liu and Hangrui Bi and Yifan " + "Li and Chun\n" + " Cai and Chengqian Zhang and Yiming Du and Jia-Xin Zhu and " + "Pinghui Mo\n" + " and Zhengtao Huang and Qiyu Zeng and Shaochen Shi and " + "Xuejian Qin and\n" + " Zhaoxi Yu and Chenxing Luo and Ye Ding and Yun-Pei Liu and " + "Ruosong Shi\n" + " and Zhenyu Wang and Sigbj{\\o}rn L{\\o}land Bore and Junhan " + "Chang and\n" + " Zhe Deng and Zhaohan Ding and Siyuan Han and Wanrun Jiang " + "and Guolin\n" + " Ke and Zhaoqing Liu and Denghui Lu and Koki Muraoka and " + "Hananeh Oliaei\n" + " and Anurag Kumar Singh and Haohui Que and Weihong Xu and " + "Zhangmancang\n" + " Xu and Yong-Bin Zhuang and Jiayu Dai and Timothy J. Giese " + "and Weile\n" + " Jia and Ben Xu and Darrin M. York and Linfeng Zhang and Han " + "Wang},\n" + " title = {{DeePMD-kit v3: A Multiple-Backend Framework for Machine " + "Learning\n" + " Potentials}},\n" + " journal = {J. Chem. Theory Comput.},\n" + " year = 2025,\n" + " volume = 21,\n" + " number = 9,\n" + " pages = {4375--4385},\n" + " doi = {10.1021/acs.jctc.5c00340},\n" + "}\n\n"; + +namespace { +// Reduced Planck constant in eV.ps. The model reports the magnetic force as +// the energy gradient with respect to the magnetic moment, while LAMMPS stores +// the precession force; the two differ by the factor hbar / |m|. +constexpr double kHBar = 6.5821191e-04; +// Positions of the LAMMPS global virial components (xx, yy, zz, xy, xz, yz) +// within the nine-component tensor the model reports. +constexpr int kGlobalVirialMap[6] = {0, 4, 8, 3, 6, 7}; +// Positions of the LAMMPS centroid per-atom virial components +// (xx, yy, zz, xy, xz, yz, yx, zx, zy) within the same tensor. +constexpr int kCentroidVirialMap[9] = {0, 4, 8, 3, 6, 7, 1, 2, 5}; +} // namespace + +PairDPA4Spin::PairDPA4Spin(LAMMPS* lmp) + : Pair(lmp), scale(nullptr), cutoff(0.0), commdata_(nullptr) { + if (lmp->citeme) { + lmp->citeme->add(cite_user_deepmd_package); + } + if (strcmp(update->unit_style, "lj") == 0) { + error->all(FLERR, + "pair style dpa4spin does not support unit style lj; use a " + "physical unit style such as metal or real."); + } + ener_unit_cvt_factor = force->boltz / 8.617343e-5; + dist_unit_cvt_factor = force->angstrom; + force_unit_cvt_factor = ener_unit_cvt_factor / dist_unit_cvt_factor; + + // The artifact is identified by a path that a restart cannot carry, so the + // input has to re-issue pair_style and pair_coeff after read_restart. + restartinfo = 0; + // The model reports a nine-component atomic virial, which feeds compute + // centroid/stress/atom. + centroidstressflag = CENTROID_AVAIL; + respa_enable = 0; + writedata = 0; + + print_summary(" "); +} + +PairDPA4Spin::~PairDPA4Spin() { + if (allocated) { + memory->destroy(setflag); + memory->destroy(cutsq); + memory->destroy(scale); + } +} + +void PairDPA4Spin::print_summary(const std::string& pre) const { + if (comm->me != 0) { + return; + } + // The DeePMD-kit banner is written to std::cout by the library. Capture it + // so that the whole summary reaches the LAMMPS screen and log file together. + std::stringstream buffer; + std::streambuf* sbuf = std::cout.rdbuf(); + std::cout.rdbuf(buffer.rdbuf()); + + std::cout << "Summary of lammps deepmd module ..." << std::endl; + std::cout << pre << ">>> Info of deepmd-kit:" << std::endl; + deep_spin.print_summary(pre); + std::cout << pre << ">>> Info of lammps module:" << std::endl; + std::cout << pre << "use deepmd-kit at: " << STR_DEEPMD_ROOT << std::endl; + std::cout << pre << "source: " << STR_GIT_SUMM << std::endl; + std::cout << pre << "source branch: " << STR_GIT_BRANCH << std::endl; + std::cout << pre << "source commit: " << STR_GIT_HASH << std::endl; + std::cout << pre << "source commit at: " << STR_GIT_DATE << std::endl; + std::cout << pre << "build with inc: " << STR_BACKEND_INCLUDE_DIRS + << std::endl; + std::cout << pre << "build with lib: " << STR_BACKEND_LIBRARY_PATH + << std::endl; + + std::cout.rdbuf(sbuf); + utils::logmesg(lmp, buffer.str()); +} + +int PairDPA4Spin::get_node_rank() const { + int rank = 0; + MPI_Comm_rank(world, &rank); + +#ifdef MPI_COMM_TYPE_SHARED + // LAMMPS may run on a partition or on an embedding-provided subcommunicator. + // Splitting that communicator by shared-memory domain keeps independent + // LAMMPS instances out of collectives on MPI_COMM_WORLD. + MPI_Comm node_comm; + MPI_Comm_split_type(world, MPI_COMM_TYPE_SHARED, rank, MPI_INFO_NULL, + &node_comm); + int node_rank = 0; + MPI_Comm_rank(node_comm, &node_rank); + MPI_Comm_free(&node_comm); + return node_rank; +#else + // The serial MPI stubs of LAMMPS predate MPI-3 and provide no + // MPI_Comm_split_type. Their only communicator holds a single rank, so its + // communicator rank is also the node-local rank. + return rank; +#endif +} + +void PairDPA4Spin::allocate() { + allocated = 1; + const int ntypes = atom->ntypes; + + memory->create(setflag, ntypes + 1, ntypes + 1, "pair:setflag"); + memory->create(cutsq, ntypes + 1, ntypes + 1, "pair:cutsq"); + memory->create(scale, ntypes + 1, ntypes + 1, "pair:scale"); + + for (int ii = 1; ii <= ntypes; ++ii) { + for (int jj = ii; jj <= ntypes; ++jj) { + setflag[ii][jj] = 0; + scale[ii][jj] = 0.0; + } + } +} + +void PairDPA4Spin::settings(int narg, char** arg) { + // Name whichever style the input selected, so the Kokkos variant reports + // itself rather than its host base. + const std::string style = force->pair_style; + if (narg != 1) { + error->all(FLERR, "Illegal pair_style command: pair style " + style + + " evaluates a single native-spin artifact and takes " + "its path as the only argument."); + } + + try { + deep_spin.init(arg[0], get_node_rank()); + } catch (deepmd_compat::deepmd_exception& e) { + error->one(FLERR, e.what()); + } + cutoff = deep_spin.cutoff() * dist_unit_cvt_factor; + + utils::logmesg(lmp, + " >>> Info of model(s):\n" + " using 1 model(s): {}\n" + " rcut in model: {}\n" + " ntypes in model: {}\n", + arg[0], cutoff, deep_spin.numb_types()); +} + +/* ---------------------------------------------------------------------- + map the atom types onto the elements of the model +------------------------------------------------------------------------- */ + +void PairDPA4Spin::coeff(int narg, char** arg) { + if (!allocated) { + allocate(); + } + + const int ntypes = atom->ntypes; + int ilo, ihi, jlo, jhi; + utils::bounds(FLERR, arg[0], 1, ntypes, ilo, ihi, error); + utils::bounds(FLERR, arg[1], 1, ntypes, jlo, jhi, error); + if (ilo != 1 || jlo != 1 || ihi != ntypes || jhi != ntypes) { + error->all(FLERR, + "pair style dpa4spin sets one scale for every atom type, i.e. " + "pair_coeff * *."); + } + + // Element names the artifact was trained on, in model type order. + std::vector model_types; + std::string type_map_str; + deep_spin.get_type_map(type_map_str); + std::istringstream type_map_stream(type_map_str); + std::string element; + while (type_map_stream >> element) { + model_types.push_back(element); + } + const int model_ntypes = static_cast(model_types.size()); + + type_idx_map.assign(ntypes, -1); + if (narg == 2) { + // A bare `pair_coeff * *` maps the atom types onto the leading model + // elements by position, which is only meaningful when the model has at + // least as many elements as the system has atom types. + if (model_ntypes < ntypes) { + error->all(FLERR, + "pair_coeff * * maps atom type i onto model element i, but " + "the system has " + + std::to_string(ntypes) + + " atom types and the model only " + + std::to_string(model_ntypes) + + "; list the elements explicitly, e.g. pair_coeff * * Fe " + "C."); + } + if (model_ntypes > ntypes) { + error->warning( + FLERR, "pair_coeff * * maps the system atom types onto the first " + + std::to_string(ntypes) + " of the model's " + + std::to_string(model_ntypes) + + " element types; list the elements explicitly, e.g. " + "pair_coeff * * Fe C, to avoid a possible mislabeling."); + } + for (int ii = 0; ii < ntypes; ++ii) { + type_idx_map[ii] = ii; + } + } else { + // An explicit element list names the model element of each atom type in + // turn. NULL, and any atom type past the end of the list, denotes a type + // the model never sees. + if (narg - 2 > ntypes) { + error->all(FLERR, + "pair_coeff lists more elements than the system has atom " + "types."); + } + for (int ii = 0; ii + 2 < narg; ++ii) { + const std::string name = arg[ii + 2]; + if (name == "NULL") { + continue; + } + const auto found = + std::find(model_types.begin(), model_types.end(), name); + if (found == model_types.end()) { + error->all(FLERR, + "pair_coeff: element " + name + " not found in the model"); + } + type_idx_map[ii] = static_cast(found - model_types.begin()); + } + } + + std::string excluded; + for (int ii = 0; ii < ntypes; ++ii) { + if (type_idx_map[ii] < 0) { + excluded += " " + std::to_string(ii + 1); + } + } + if (!excluded.empty()) { + error->warning(FLERR, "pair style dpa4spin ignores atom type(s)" + + excluded + ": they map to no model element."); + } + + for (int ii = 1; ii <= ntypes; ++ii) { + for (int jj = ii; jj <= ntypes; ++jj) { + setflag[ii][jj] = 1; + scale[ii][jj] = 1.0; + } + } +} + +void PairDPA4Spin::init_style() { + neighbor->add_request(this, NeighConst::REQ_FULL); + + const std::string style = force->pair_style; + if (!atom->sp_flag) { + error->all(FLERR, "pair style " + style + + " only supports spin atoms, please use pair style " + "deepmd instead."); + } + // The scheme, not the lower schema, decides which style serves an artifact: + // this one marshals one node per atom and reads the magnetic force straight + // off the model, neither of which a virtual-atom model provides. + if (!deep_spin.uses_native_spin_scheme()) { + error->all(FLERR, + "pair style " + style + + " serves the native-spin scheme without cross-rank message " + "passing; a virtual-atom spin model, and a native-spin " + "model whose descriptor exchanges intermediate features " + "between ranks, are served by pair style deepspin."); + } + // A single rank folds ghost neighbours onto the local atom that owns them, + // which is resolved through the atom map. Domain decomposition gives every + // ghost its own node and needs no map. + if (comm->nprocs == 1 && atom->map_style == Atom::MAP_NONE) { + error->all(FLERR, "pair style " + style + + " needs an atom map on a single rank; add " + "'atom_modify map yes' to the input."); + } +} + +double PairDPA4Spin::init_one(int i, int j) { + if (setflag[i][j] == 0) { + scale[i][j] = 1.0; + } + scale[j][i] = scale[i][j]; + + return cutoff; +} + +void PairDPA4Spin::compute(int eflag, int vflag) { + ev_init(eflag, vflag); + if (vflag_atom) { + error->all(FLERR, + "6-element atomic virial is not supported. Use compute " + "centroid/stress/atom command for 9-element atomic virial."); + } + + const int nlocal = atom->nlocal; + const int nghost = atom->nghost; + const int nall = nlocal + nghost; + if (nall == 0) { + // A rank holding no atom at all contributes nothing to any output. + return; + } + + double** x = atom->x; + double** f = atom->f; + double** sp = atom->sp; + double** fm = atom->fm; + const int* type = atom->type; + + // === Step 1. Marshal the model inputs === + // Coordinates are referred to the box origin, and the moment is the LAMMPS + // unit direction scaled by its magnitude. + std::vector dcoord(static_cast(nall) * 3); + std::vector dspin(static_cast(nall) * 3); + std::vector dtype(nall); + for (int ii = 0; ii < nall; ++ii) { + for (int dd = 0; dd < 3; ++dd) { + dcoord[ii * 3 + dd] = + (x[ii][dd] - domain->boxlo[dd]) / dist_unit_cvt_factor; + dspin[ii * 3 + dd] = sp[ii][dd] * sp[ii][3]; + } + dtype[ii] = type_idx_map[type[ii] - 1]; + } + + std::vector dbox(9, 0.0); + dbox[0] = domain->h[0] / dist_unit_cvt_factor; // xx + dbox[4] = domain->h[1] / dist_unit_cvt_factor; // yy + dbox[8] = domain->h[2] / dist_unit_cvt_factor; // zz + dbox[7] = domain->h[3] / dist_unit_cvt_factor; // zy + dbox[6] = domain->h[4] / dist_unit_cvt_factor; // zx + dbox[3] = domain->h[5] / dist_unit_cvt_factor; // yx + + commdata_ = (CommBrickDPA4Spin*)comm; + deepmd_compat::InputNlist lmp_list( + list->inum, list->ilist, list->numneigh, list->firstneigh, + commdata_->nswap, commdata_->sendnum, commdata_->recvnum, + commdata_->firstrecv, commdata_->sendlist, commdata_->sendproc, + commdata_->recvproc, &world, comm->nprocs); + lmp_list.set_mask(NEIGHMASK); + // A single rank folds every ghost onto the local atom that owns it, an + // ownership the atom map resolves. Domain decomposition gives each ghost a + // node of its own and needs no such map. + std::vector mapping; + if (comm->nprocs == 1) { + mapping.resize(nall); + for (int ii = 0; ii < nall; ++ii) { + mapping[ii] = atom->map(atom->tag[ii]); + } + lmp_list.set_mapping(mapping.data()); + } + + // === Step 2. Evaluate the model === + // LAMMPS resets ago to zero on every neighbor-list rebuild, which is when + // the backend refreshes its cached topology. + double dener = 0.0; + std::vector dforce(static_cast(nall) * 3); + std::vector dforce_mag(static_cast(nall) * 3); + std::vector dvirial(9, 0.0); + std::vector deatom, dvatom; + try { + if (eflag_atom || cvflag_atom) { + deep_spin.compute(dener, dforce, dforce_mag, dvirial, deatom, dvatom, + dcoord, dspin, dtype, dbox, nghost, lmp_list, + neighbor->ago); + } else { + deep_spin.compute(dener, dforce, dforce_mag, dvirial, dcoord, dspin, + dtype, dbox, nghost, lmp_list, neighbor->ago); + } + } catch (deepmd_compat::deepmd_exception& e) { + error->one(FLERR, e.what()); + } + + // === Step 3. Accumulate the model outputs === + // The model reports a force and a magnetic force for every atom the neighbor + // list covers, ghosts included; the spin atom style folds the ghost rows + // onto their owners in its reverse communication. Dividing the magnetic + // force by hbar / |m| turns the energy gradient with respect to the moment + // into the precession force LAMMPS stores, and leaves an atom whose moment + // vanishes with no precession force at all. + for (int ii = 0; ii < nall; ++ii) { + for (int dd = 0; dd < 3; ++dd) { + f[ii][dd] += scale[1][1] * dforce[3 * ii + dd] * force_unit_cvt_factor; + fm[ii][dd] += scale[1][1] * dforce_mag[3 * ii + dd] / + (kHBar / sp[ii][3]) * force_unit_cvt_factor; + } + } + + if (eflag) { + eng_vdwl += scale[1][1] * dener * ener_unit_cvt_factor; + } + if (vflag) { + for (int kk = 0; kk < 6; ++kk) { + virial[kk] += + scale[1][1] * dvirial[kGlobalVirialMap[kk]] * ener_unit_cvt_factor; + } + } + if (eflag_atom) { + for (int ii = 0; ii < nlocal; ++ii) { + eatom[ii] += scale[1][1] * deatom[ii] * ener_unit_cvt_factor; + } + } + if (cvflag_atom) { + for (int ii = 0; ii < nall; ++ii) { + for (int kk = 0; kk < 9; ++kk) { + cvatom[ii][kk] += scale[1][1] * + dvatom[9 * ii + kCentroidVirialMap[kk]] * + ener_unit_cvt_factor; + } + } + } +} + +void* PairDPA4Spin::extract(const char* str, int& dim) { + if (strcmp(str, "scale") == 0) { + dim = 2; + return (void*)scale; + } + return nullptr; +} diff --git a/source/lmp/pair_dpa4spin.h b/source/lmp/pair_dpa4spin.h new file mode 100644 index 0000000000..73ac544734 --- /dev/null +++ b/source/lmp/pair_dpa4spin.h @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#ifndef LAMMPS_VERSION_NUMBER +#error Please define LAMMPS_VERSION_NUMBER to yyyymmdd +#endif + +#ifdef PAIR_CLASS +// clang-format off +PairStyle(dpa4spin,PairDPA4Spin); +// clang-format on +#else + +#ifndef LMP_PAIR_DPA4SPIN_H +#define LMP_PAIR_DPA4SPIN_H + +#ifdef DP_USE_CXX_API +#ifdef LMPPLUGIN +#include "DeepSpin.h" +#else +#include "deepmd/DeepSpin.h" +#endif +namespace deepmd_compat = deepmd; +#else +#ifdef LMPPLUGIN +#include "deepmd.hpp" +#else +#include "deepmd/deepmd.hpp" +#endif +namespace deepmd_compat = deepmd::hpp; +#endif + +#include +#include + +#include "comm_brick.h" +#include "pair.h" + +namespace LAMMPS_NS { + +// Opens the ghost-swap metadata that CommBrick keeps protected. The host +// neighbor-list interface of the model carries it so that a domain-decomposed +// run can describe its halo to the backend. +class CommBrickDPA4Spin : public CommBrick { + friend class PairDPA4Spin; +}; + +// Host pair style for native-spin models, where the magnetic moment enters the +// descriptor as an equivariant input and there are no virtual atoms. Either +// graph lower the scheme defines is served; the artifact declares which one +// and the backend branches on it. +// +// Coordinates and moments are marshaled to the host neighbor-list interface of +// the model, which returns the per-atom force, magnetic force, energy and +// virial. Under domain decomposition the model reports extended (local plus +// ghost) force and magnetic force, both of which the spin atom style folds +// onto their owners through its reverse communication. +// +// The style evaluates exactly one artifact and passes no frame, atomic or +// charge/spin parameters, so a model that requires one must carry its default. +// +// The device-resident variant is ``dpa4spin/kk``; it needs the compact +// canonical artifact and evaluates it without the per-step host marshaling. +class PairDPA4Spin : public Pair { + public: + PairDPA4Spin(class LAMMPS*); + ~PairDPA4Spin() override; + + // Load the artifact named by ``pair_style dpa4spin ``. + void settings(int, char**) override; + // Resolve the LAMMPS atom types onto the element list of the model. + void coeff(int, char**) override; + // Request the neighbor list and reject every setting the scheme cannot + // serve. The Kokkos variant chains through this method and then adds the + // requirements of device residency. + void init_style() override; + double init_one(int, int) override; + void compute(int, int) override; + void* extract(const char*, int&) override; + + protected: + void allocate(); + // Emit the DeePMD-kit and LAMMPS module banner through the LAMMPS logger. + void print_summary(const std::string& pre) const; + // Rank within the shared-memory domain, which selects the accelerator the + // model binds to. + int get_node_rank() const; + + deepmd_compat::DeepSpin deep_spin; + // Per-type-pair prefactor applied to every model output; ``fix adapt`` + // reaches it through extract(). + double** scale; + // Interaction cutoff of the model, in LAMMPS distance units. + double cutoff; + // LAMMPS atom type (zero based) -> model element index, or -1 for a type the + // model never sees. + std::vector type_idx_map; + // LAMMPS unit system relative to the eV / Angstrom system of the model. + double ener_unit_cvt_factor, dist_unit_cvt_factor, force_unit_cvt_factor; + + private: + CommBrickDPA4Spin* commdata_; +}; + +} // namespace LAMMPS_NS + +#endif +#endif diff --git a/source/lmp/pair_dpa4spin_kokkos.cpp b/source/lmp/pair_dpa4spin_kokkos.cpp new file mode 100644 index 0000000000..8ae618e752 --- /dev/null +++ b/source/lmp/pair_dpa4spin_kokkos.cpp @@ -0,0 +1,475 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#ifdef LMP_KOKKOS +#include "pair_dpa4spin_kokkos.h" + +#include + +#include "atom.h" +#include "atom_kokkos.h" +#include "atom_masks.h" +#include "comm.h" +#include "error.h" +#include "kokkos.h" +#include "memory_kokkos.h" +#include "neigh_list_kokkos.h" +#include "neigh_request.h" +#include "neighbor.h" + +using namespace LAMMPS_NS; + +namespace { +// Reduced Planck constant in eV.ps. The model reports the magnetic force as +// the energy gradient with respect to the magnetic moment, while LAMMPS stores +// the precession force; the two differ by the moment magnitude over hbar. +constexpr double kHBar = 6.5821191e-04; +} // namespace + +template +PairDPA4SpinKokkos::PairDPA4SpinKokkos(LAMMPS* lmp) + : PairDPA4Spin(lmp), + compact_graph(lmp), + multi_rank(false), + reverse_virial(false), + reverse_used_host(false) { + respa_enable = 0; + kokkosable = 1; + atomKK = (AtomKokkos*)atom; + execution_space = ExecutionSpaceFromDevice::space; + datamask_read = X_MASK | TYPE_MASK | SP_MASK | ENERGY_MASK | VIRIAL_MASK; + datamask_modify = F_MASK | FM_MASK | ENERGY_MASK | VIRIAL_MASK; + reverse_comm_device = 1; +} + +template +PairDPA4SpinKokkos::~PairDPA4SpinKokkos() { + if (copymode) { + return; + } + memoryKK->destroy_kokkos(k_eatom, eatom); +} + +template +int PairDPA4SpinKokkos::pack_reverse_comm(int n, + int first, + double* buf) { + if (reverse_virial) { + auto h_reverse = k_reverse_virial.view_host(); + int m = 0; + const int last = first + n; + for (int i = first; i < last; ++i) { + for (int k = 0; k < 9; ++k) { + buf[m++] = h_reverse(9 * i + k); + } + } + return m; + } + reverse_used_host = true; + atomKK->sync(Host, F_MASK | FM_MASK); + double** f = atom->f; + double** fm = atom->fm; + int m = 0; + const int last = first + n; + for (int i = first; i < last; ++i) { + buf[m++] = f[i][0]; + buf[m++] = f[i][1]; + buf[m++] = f[i][2]; + buf[m++] = fm[i][0]; + buf[m++] = fm[i][1]; + buf[m++] = fm[i][2]; + } + return m; +} + +template +void PairDPA4SpinKokkos::unpack_reverse_comm(int n, + int* list, + double* buf) { + if (reverse_virial) { + k_reverse_virial.modify_host(); + auto h_reverse = k_reverse_virial.view_host(); + int m = 0; + for (int i = 0; i < n; ++i) { + const int j = list[i]; + for (int k = 0; k < 9; ++k) { + h_reverse(9 * j + k) += buf[m++]; + } + } + return; + } + reverse_used_host = true; + atomKK->sync(Host, F_MASK | FM_MASK); + double** f = atom->f; + double** fm = atom->fm; + int m = 0; + for (int i = 0; i < n; ++i) { + const int j = list[i]; + f[j][0] += buf[m++]; + f[j][1] += buf[m++]; + f[j][2] += buf[m++]; + fm[j][0] += buf[m++]; + fm[j][1] += buf[m++]; + fm[j][2] += buf[m++]; + } + atomKK->modified(Host, F_MASK | FM_MASK); +} + +template +int PairDPA4SpinKokkos::pack_reverse_comm_kokkos( + int n, int first, DAT::tdual_double_1d& buf) { + auto d_buf = buf.template view(); + if (reverse_virial) { + auto reverse_virial_data = k_reverse_virial.template view(); + const int first_i = first; + Kokkos::parallel_for( + "dpa4spin/kk:pack_rev_virial", Kokkos::RangePolicy(0, n), + KOKKOS_LAMBDA(const int i) { + for (int k = 0; k < 9; ++k) { + d_buf(9 * i + k) = reverse_virial_data(9 * (first_i + i) + k); + } + }); + return n * 9; + } + auto f = atomKK->k_f.template view(); + auto fm = atomKK->k_fm.template view(); + const int first_i = first; + Kokkos::parallel_for( + "dpa4spin/kk:pack_rev", Kokkos::RangePolicy(0, n), + KOKKOS_LAMBDA(const int i) { + d_buf(6 * i + 0) = f(first_i + i, 0); + d_buf(6 * i + 1) = f(first_i + i, 1); + d_buf(6 * i + 2) = f(first_i + i, 2); + d_buf(6 * i + 3) = fm(first_i + i, 0); + d_buf(6 * i + 4) = fm(first_i + i, 1); + d_buf(6 * i + 5) = fm(first_i + i, 2); + }); + return n * 6; +} + +template +void PairDPA4SpinKokkos::unpack_reverse_comm_kokkos( + int n, DAT::tdual_int_1d list, DAT::tdual_double_1d& buf) { + auto d_buf = buf.template view(); + auto d_list = list.template view(); + if (reverse_virial) { + k_reverse_virial.template modify(); + auto reverse_virial_data = k_reverse_virial.template view(); + Kokkos::parallel_for( + "dpa4spin/kk:unpack_rev_virial", Kokkos::RangePolicy(0, n), + KOKKOS_LAMBDA(const int i) { + const int j = d_list(i); + for (int k = 0; k < 9; ++k) { + reverse_virial_data(9 * j + k) += d_buf(9 * i + k); + } + }); + return; + } + auto f = atomKK->k_f.template view(); + auto fm = atomKK->k_fm.template view(); + Kokkos::parallel_for( + "dpa4spin/kk:unpack_rev", Kokkos::RangePolicy(0, n), + KOKKOS_LAMBDA(const int i) { + const int j = d_list(i); + f(j, 0) += d_buf(6 * i + 0); + f(j, 1) += d_buf(6 * i + 1); + f(j, 2) += d_buf(6 * i + 2); + fm(j, 0) += d_buf(6 * i + 3); + fm(j, 1) += d_buf(6 * i + 4); + fm(j, 2) += d_buf(6 * i + 5); + }); +} + +template +void PairDPA4SpinKokkos::init_style() { + // Full neighbor-list request and the shared native-spin contract. + PairDPA4Spin::init_style(); + + // The device-resident entry point runs the model on an accelerator, so the + // Kokkos host execution space has nothing to hand it. + if (std::is_same::value) { + error->all(FLERR, "pair style dpa4spin/kk runs on the GPU backend only."); + } + // Device residency is what the compact canonical ABI exists for: it is the + // only schema whose inputs this style can hand over without a host round + // trip. + if (!deep_spin.uses_canonical_graph_inference()) { + error->all(FLERR, + "pair style dpa4spin/kk requires a model frozen with the " + "compact canonical graph lower; a model frozen with the graph " + "lower is served by pair style dpa4spin."); + } + // Domain decomposition uses the extended (local + ghost) node set: the model + // computes per-node force and magnetic force and the reverse communication + // folds the ghost contributions onto their owners. A single rank uses the + // folded minimum-image node set. + multi_rank = (comm->nprocs > 1); + + // Route the base full request to the Kokkos device neighbor build. + auto request = neighbor->find_request(this); + request->set_kokkos_device(std::is_same::value); + request->set_kokkos_host(false); + request->enable_full(); + // Force and magnetic force exchange three values each per atom. Centroid + // per-atom virial uses nine values even though the Kokkos full-list request + // runs newton off, so comm_reverse_off reserves the classic host buffer for + // the wider mode. + comm_reverse = 6; + comm_reverse_off = 9; + + // The model node set is shared with the graph builder; type_idx_map is + // populated by the base coeff(). + compact_graph.setup(type_idx_map, multi_rank); +} + +template +void PairDPA4SpinKokkos::gather_moment_device() { + const int nall = atom->nlocal + atom->nghost; + const int nnode_model = compact_graph.nnode_model; + if ((int)d_moment.extent(0) < 3 * nnode_model) { + d_moment = Kokkos::View("dpa4spin/kk:moment", 3 * nall); + } + // LAMMPS stores the moment as a unit direction and its magnitude; the + // artifact consumes the product. Ghost rows repeat their owner's value, + // which the ``sp`` forward communication has already refreshed, so the + // moment needs no exchange of its own. + atomKK->sync(execution_space, SP_MASK); + auto sp = atomKK->k_sp.template view(); + auto moment = d_moment; + auto model2loc = compact_graph.d_model2loc; + Kokkos::parallel_for( + "dpa4spin/kk:moment", Kokkos::RangePolicy(0, nnode_model), + KOKKOS_LAMBDA(const int m) { + const int i = model2loc(m); + const auto norm = sp(i, 3); + moment(3 * m + 0) = static_cast(sp(i, 0) * norm); + moment(3 * m + 1) = static_cast(sp(i, 1) * norm); + moment(3 * m + 2) = static_cast(sp(i, 2) * norm); + }); +} + +template +void PairDPA4SpinKokkos::compute(int eflag, int vflag) { + ev_init(eflag, vflag); + if (vflag_atom) { + error->all(FLERR, + "6-element atomic virial is not supported. Use compute " + "centroid/stress/atom command for 9-element atomic virial."); + } + + const int nlocal = atom->nlocal; + // Per-atom energy is scattered on the device into a DualView that aliases the + // base Pair ``eatom`` array; (re)allocate it here as the standard Kokkos + // pair styles do. The centroid per-atom virial has no Kokkos device path, so + // it is filled on the host below. + if (eflag_atom) { + memoryKK->destroy_kokkos(k_eatom, eatom); + memoryKK->create_kokkos(k_eatom, eatom, maxeatom, "dpa4spin/kk:eatom"); + d_eatom = k_eatom.template view(); + } + compact_graph.build(list, cutoff, dist_unit_cvt_factor); + gather_moment_device(); + const int nloc_m = compact_graph.nloc_model; // local nodes (energy) + const int nnode_m = compact_graph.nnode_model; // all nodes (force / virial) + + // Energy is per local node; force, magnetic force and virial span the model + // node set, which is the local atoms (folded) or local + real ghost atoms + // (extended, up to nall). The two extents grow independently: under domain + // decomposition ``nlocal`` and ``nall`` need not move together, so a shared + // guard could leave the energy buffer short when ``nlocal`` grows while + // ``nall`` does not. + const int nall = atom->nlocal + atom->nghost; + if ((int)d_atom_energy.extent(0) < nlocal) { + d_atom_energy = + Kokkos::View("dpa4spin/kk:atom_energy", nlocal); + } + if ((int)d_out_force.extent(0) < 3 * nall) { + d_out_force = + Kokkos::View("dpa4spin/kk:out_force", 3 * nall); + d_out_force_mag = Kokkos::View( + "dpa4spin/kk:out_force_mag", 3 * nall); + d_atom_virial = + Kokkos::View("dpa4spin/kk:atom_virial", 9 * nall); + } + if (cvflag_atom && multi_rank && (int)k_reverse_virial.extent(0) < 9 * nall) { + k_reverse_virial = + DAT::tdual_double_1d("dpa4spin/kk:reverse_virial", 9 * nall); + } + Kokkos::deep_copy(d_out_force, 0.0); + Kokkos::deep_copy(d_out_force_mag, 0.0); + Kokkos::deep_copy(d_atom_energy, 0.0); + Kokkos::deep_copy(d_atom_virial, 0.0); + + if (nnode_m > 0) { + // Fully device-resident inference: raw device pointers in and out. The + // graph and the moment are produced on the Kokkos stream and consumed by + // the model on PyTorch's stream, and the outputs flow back to the Kokkos + // scatter, so the two runtimes are bracketed by explicit synchronization: + // fence the Kokkos work before the model reads its inputs, and synchronize + // the device after so the scatter sees the finished model outputs. + Kokkos::fence(); + try { + deep_spin.compute_canonical_graph_gpu( + d_atom_energy.data(), d_out_force.data(), d_out_force_mag.data(), + d_atom_virial.data(), compact_graph.d_model_type.data(), + compact_graph.d_source.data(), compact_graph.d_edge_vec.data(), + compact_graph.d_destination_row_ptr.data(), + compact_graph.d_source_row_ptr.data(), + compact_graph.d_source_order.data(), d_moment.data(), nloc_m, nnode_m, + compact_graph.storage_count); + } catch (deepmd_compat::deepmd_exception& e) { + error->one(FLERR, e.what()); + } + } + + // === Scatter the model-node forces onto their atoms === + // ``model2loc`` maps a model node back to its LAMMPS atom (the identity when + // there are no virtual atoms); virtual atoms receive no contribution. For the + // extended multi-domain set the nodes past ``nloc_m`` are ghosts, whose + // forces are written to the ghost slots and folded onto their owners by the + // reverse communication that the KOKKOS package (which forces 'newton off' + // with a full list) would otherwise skip. + // The scatter remains device-resident. If LAMMPS selects classic host + // communication, the host pack/unpack methods synchronize the force + // DualViews and the completed fold is copied back once after all + // communication stages. + auto model2loc = compact_graph.d_model2loc; + const double fscale = scale[1][1] * force_unit_cvt_factor; + const double fmscale = scale[1][1] * force_unit_cvt_factor / kHBar; + reverse_virial = false; + reverse_used_host = false; + // The KOKKOS package runs 'newton off', so the integrator's force_clear only + // zeros the local force and magnetic force (indices [0, nlocal)); the ghost + // slots [nlocal, nall) are left untouched. The extended scatter writes ghost + // slots and folds them onto their owners by reverse communication, so those + // slots must be zeroed first, or their contribution accumulates across steps. + atomKK->sync(execution_space, F_MASK | FM_MASK | SP_MASK); + auto f = atomKK->k_f.template view(); + auto fm = atomKK->k_fm.template view(); + auto sp = atomKK->k_sp.template view(); + auto out_force = d_out_force; + auto out_force_mag = d_out_force_mag; + if (multi_rank) { + Kokkos::parallel_for( + "dpa4spin/kk:clear_ghost_f", + Kokkos::RangePolicy(nlocal, nall), + KOKKOS_LAMBDA(const int m) { + f(m, 0) = 0.0; + f(m, 1) = 0.0; + f(m, 2) = 0.0; + fm(m, 0) = 0.0; + fm(m, 1) = 0.0; + fm(m, 2) = 0.0; + }); + } + Kokkos::parallel_for( + "dpa4spin/kk:scatter_f", Kokkos::RangePolicy(0, nnode_m), + KOKKOS_LAMBDA(const int m) { + const int i = model2loc(m); + // The precession force carried by ``fm`` is the energy gradient with + // respect to the moment, rescaled by the moment magnitude over hbar. + const double moment = fmscale * sp(i, 3); + f(i, 0) += fscale * out_force(3 * m + 0); + f(i, 1) += fscale * out_force(3 * m + 1); + f(i, 2) += fscale * out_force(3 * m + 2); + fm(i, 0) += moment * out_force_mag(3 * m + 0); + fm(i, 1) += moment * out_force_mag(3 * m + 1); + fm(i, 2) += moment * out_force_mag(3 * m + 2); + }); + atomKK->modified(execution_space, F_MASK | FM_MASK); + if (multi_rank) { + comm->reverse_comm(this, 6); + if (reverse_used_host) { + atomKK->sync(execution_space, F_MASK | FM_MASK); + } + } + + if (eflag_global) { + auto atom_energy = d_atom_energy; + double e_sum = 0.0; + Kokkos::parallel_reduce( + "dpa4spin/kk:esum", Kokkos::RangePolicy(0, nloc_m), + KOKKOS_LAMBDA(const int m, double& acc) { acc += atom_energy(m); }, + e_sum); + eng_vdwl += scale[1][1] * e_sum * ener_unit_cvt_factor; + } + + if (vflag_global) { + // Sum the per-node 9-component virial and map to the LAMMPS global 6 + // (xx, yy, zz, xy, xz, yz), matching the standalone pair's index map. The + // sum spans all nodes (local + extended ghost) so the reduction equals the + // model's reduced virial for this rank's local-centered edges. + auto atom_virial = d_atom_virial; + const int comp[6] = {0, 4, 8, 3, 6, 7}; + for (int k = 0; k < 6; ++k) { + const int off = comp[k]; + double vsum = 0.0; + Kokkos::parallel_reduce( + "dpa4spin/kk:vsum", Kokkos::RangePolicy(0, nnode_m), + KOKKOS_LAMBDA(const int m, double& acc) { + acc += atom_virial(9 * m + off); + }, + vsum); + virial[k] += scale[1][1] * vsum * ener_unit_cvt_factor; + } + } + + if (eflag_atom) { + auto atom_energy = d_atom_energy; + auto eatom_v = d_eatom; + const double escale = scale[1][1] * ener_unit_cvt_factor; + Kokkos::deep_copy(d_eatom, 0.0); // virtual atoms keep zero energy + Kokkos::parallel_for( + "dpa4spin/kk:eatom", Kokkos::RangePolicy(0, nloc_m), + KOKKOS_LAMBDA(const int m) { + eatom_v(model2loc(m)) = escale * atom_energy(m); + }); + k_eatom.template modify(); + k_eatom.sync_host(); + } + + if (cvflag_atom) { + // Centroid per-atom virial is reported on owned atoms. Contributions + // carried by extended ghost nodes are folded to their owners explicitly + // because the KOKKOS full-list path runs with newton pair disabled. + auto h_av = Kokkos::create_mirror_view(d_atom_virial); + Kokkos::deep_copy(h_av, d_atom_virial); + auto h_m2l = compact_graph.k_model2loc.view_host(); + const double vscale = scale[1][1] * ener_unit_cvt_factor; + const int map9[9] = {0, 4, 8, 3, 6, 7, 1, 2, 5}; + for (int m = 0; m < nloc_m; ++m) { + const int ii = h_m2l(m); + for (int k = 0; k < 9; ++k) { + cvatom[ii][k] += vscale * h_av(9 * m + map9[k]); + } + } + if (multi_rank) { + reverse_virial = true; + k_reverse_virial.modify_host(); + auto h_reverse = k_reverse_virial.view_host(); + Kokkos::deep_copy(h_reverse, 0.0); + for (int m = nloc_m; m < nnode_m; ++m) { + const int ii = h_m2l(m); + for (int k = 0; k < 9; ++k) { + h_reverse(9 * ii + k) = vscale * h_av(9 * m + map9[k]); + } + } + k_reverse_virial.template sync(); + comm->reverse_comm(this, 9); + k_reverse_virial.sync_host(); + for (int i = 0; i < nlocal; ++i) { + for (int k = 0; k < 9; ++k) { + cvatom[i][k] += h_reverse(9 * i + k); + } + } + reverse_virial = false; + } + } +} + +namespace LAMMPS_NS { +template class PairDPA4SpinKokkos; +#ifdef LMP_KOKKOS_GPU +template class PairDPA4SpinKokkos; +#endif +} // namespace LAMMPS_NS + +#endif // LMP_KOKKOS diff --git a/source/lmp/pair_dpa4spin_kokkos.h b/source/lmp/pair_dpa4spin_kokkos.h new file mode 100644 index 0000000000..d4506693f0 --- /dev/null +++ b/source/lmp/pair_dpa4spin_kokkos.h @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// The device pair style is available when the LAMMPS Kokkos package is enabled. +#ifdef LMP_KOKKOS + +#ifndef LAMMPS_VERSION_NUMBER +#error Please define LAMMPS_VERSION_NUMBER to yyyymmdd +#endif + +#ifdef PAIR_CLASS +// clang-format off +PairStyle(dpa4spin/kk,PairDPA4SpinKokkos); +PairStyle(dpa4spin/kk/device,PairDPA4SpinKokkos); +PairStyle(dpa4spin/kk/host,PairDPA4SpinKokkos); +// clang-format on +#else + +#ifndef LMP_PAIR_DPA4SPIN_KOKKOS_H +#define LMP_PAIR_DPA4SPIN_KOKKOS_H + +#include +#include + +#include "compact_canonical_graph_kokkos.h" +#include "kokkos_base.h" +#include "kokkos_type.h" +#include "neigh_list_kokkos.h" +#include "pair_dpa4spin.h" + +namespace LAMMPS_NS { + +// GPU-resident inference for exported native-spin ``.pt2`` models whose forward +// consumes the compact canonical graph: a dual-CSR neighbor topology with +// uint32 indices and float32 edge vectors, plus the per-node magnetic moment. +// It is dispatched through ``DeepSpin::compute_canonical_graph_gpu``. +// +// The neighbor list, the compact graph, the moment and the model outputs all +// stay on the device: the graph is built from the Kokkos device neighbor list, +// the moment is gathered from the Kokkos ``sp`` array, both are handed to +// ``compute_canonical_graph_gpu`` as raw device pointers, and the returned +// per-atom force, magnetic force, energy and virial are scattered back into the +// Kokkos atom arrays without any host round-trip. This removes the per-step +// host coordinate and moment marshaling of the ``dpa4spin`` path. +// +// A single rank uses the folded minimum-image node set (box thickness +// > 2 * cutoff along every periodic direction); domain decomposition uses the +// extended local-plus-ghost node set and folds ghost force and magnetic force +// onto their owners through reverse communication. +template +class PairDPA4SpinKokkos : public PairDPA4Spin, public KokkosBase { + public: + typedef DeviceType device_type; + typedef ArrayTypes AT; + + PairDPA4SpinKokkos(class LAMMPS*); + ~PairDPA4SpinKokkos() override; + + void compute(int, int) override; + void init_style() override; + // Fold extended (ghost) node outputs onto their owners. The KOKKOS package + // forces 'newton off' with a full neighbor list, disabling the integrator's + // automatic reverse communication, so the extended multi-domain path drives + // it explicitly for force, magnetic force and centroid per-atom virial. The + // Kokkos overrides run device-resident with GPU-aware MPI; the plain + // overrides serve the host-staged path. + int pack_reverse_comm(int, int, double*) override; + void unpack_reverse_comm(int, int*, double*) override; + int pack_reverse_comm_kokkos(int, int, DAT::tdual_double_1d&) override; + void unpack_reverse_comm_kokkos(int, + DAT::tdual_int_1d, + DAT::tdual_double_1d&) override; + + // Gather the per-node magnetic moment from the Kokkos ``sp`` array. Public + // because it launches an extended device lambda, which CUDA forbids inside + // non-public members. + void gather_moment_device(); + + protected: + // Model node set and the compact canonical graph over it. + CompactCanonicalGraphKokkos compact_graph; + bool multi_rank; // domain-decomposed run -> extended (local+ghost) node set + // (3 * nall) per-node magnetic moment in the ABI's float32 layout, the + // LAMMPS unit direction scaled by its magnitude. + Kokkos::View d_moment; + + // Model outputs on the device. Energy is per local atom; force, magnetic + // force and virial span the model node set (up to ``nall`` under domain + // decomposition). + Kokkos::View d_atom_energy; // (nlocal) + Kokkos::View d_out_force; // (3 * nall) + Kokkos::View d_out_force_mag; // (3 * nall) + Kokkos::View d_atom_virial; // (9 * nall) + DAT::tdual_double_1d + k_reverse_virial; // (9 * nall), atom-order ghost contributions + + // Per-atom energy accumulator (aliases the base Pair ``eatom`` host array so + // downstream per-atom computes/dumps see it after the device-to-host sync). + DAT::ttransform_kkacc_1d k_eatom; + typename AT::t_kkacc_1d d_eatom; + + bool reverse_virial; // reverse communication operates on centroid virial + bool reverse_used_host; // force reverse communication selected host staging +}; + +} // namespace LAMMPS_NS + +#endif +#endif + +#endif // LMP_KOKKOS diff --git a/source/op/pt/dpa1_graph_energy_force.cu b/source/op/pt/dpa1_graph_energy_force.cu index 212e4a7550..0c57805a47 100644 --- a/source/op/pt/dpa1_graph_energy_force.cu +++ b/source/op/pt/dpa1_graph_energy_force.cu @@ -101,8 +101,8 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, const torch::Tensor& g_saved = std::get<6>(desc); // === Step 2. Fitting forward: descriptor -> per-atom energy. === - auto fit = graph_fitting(grrg, atype, fit_ws, fit_bs, fit_resnets, - w_head, b_head, bias_atom_e, fit_act); + auto fit = graph_fitting(grrg, atype, fit_ws, fit_bs, fit_resnets, w_head, + b_head, bias_atom_e, fit_act); const torch::Tensor& atom_energy_raw = std::get<0>(fit); // (N, 1) fp64 const torch::Tensor& fit_saved = std::get<1>(fit); auto owned = ownership.reshape({-1, 1}).to(atom_energy_raw.scalar_type()); @@ -130,10 +130,11 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, // === Step 5. Scatter dE/d(edge_vec) into force / virial / atom virial. === // g_e and edge_vec_f are already in the compute precision; the per-node force // is a short neighbor sum and the per-frame virial reduces hierarchically. - auto fv = edge_force_virial(g_e, edge_vec_f, edge_index, edge_mask, - destination_order, destination_row_ptr, - source_order, source_row_ptr, n_node, - node_capacity, do_atomic_virial); + // DPA1 has no magnetic degree of freedom, so the spin cotangent is absent. + auto fv = edge_force_virial( + g_e, edge_vec_f, edge_index, edge_mask, destination_order, + destination_row_ptr, source_order, source_row_ptr, n_node, + torch::empty({0}, edge_vec_f.options()), node_capacity, do_atomic_virial); return {energy, atom_energy, std::get<0>(fv), std::get<2>(fv), std::get<1>(fv)}; } @@ -151,7 +152,8 @@ TORCH_LIBRARY_FRAGMENT(deepmd, m) { "int act, int type_one_side, int concat_tebd, int smooth, int axis, int " "resnet2, int resnet3, float rcut, float rcut_smth, float protection, " "float nnei, int basis_dim, Tensor[] fit_ws, Tensor[] fit_bs, " - "int[] fit_resnets, Tensor w_head, Tensor b_head, Tensor bias_atom_e, int " + "int[] fit_resnets, Tensor w_head, Tensor b_head, Tensor bias_atom_e, " + "int " "fit_act, SymInt node_capacity, bool do_atomic_virial) -> (Tensor, " "Tensor, Tensor, Tensor, Tensor)"); m.impl("dpa1_graph_energy_force", torch::kCUDA, &dpa1_graph_energy_force); diff --git a/source/op/pt/dpa4c_graph_compress.cu b/source/op/pt/dpa4c_graph_compress.cu index 74db36f8de..e70f96b9ae 100644 --- a/source/op/pt/dpa4c_graph_compress.cu +++ b/source/op/pt/dpa4c_graph_compress.cu @@ -38,33 +38,35 @@ struct Dimensions { int output_width; int degree_one; int coupling_records; + int spin_channels; }; -template +template Dimensions dimensions_of() { - using P = Profile; + using P = Profile; return {P::MomentWidth, P::OutputWidth, P::C1, - deepmd_dpa4c::coupling_record_count(Lmax)}; + deepmd_dpa4c::coupling_record_count(Lmax), P::Cs}; } -template +template Dimensions dimensions_for(int lmax) { switch (lmax) { case 2: - return dimensions_of(); + return dimensions_of(); case 3: - return dimensions_of(); + return dimensions_of(); case 4: - return dimensions_of(); + return dimensions_of(); default: TORCH_CHECK(false, "dpa4c_graph_compress: unsupported lmax ", lmax); } } -Dimensions profile_dimensions(int channels, int lmax) { -#define DPA4C_DIMENSIONS(width) \ - if (channels == width) { \ - return dimensions_for(lmax); \ +Dimensions profile_dimensions(int channels, int lmax, bool has_spin) { +#define DPA4C_DIMENSIONS(width) \ + if (channels == width) { \ + return has_spin ? dimensions_for(lmax) \ + : dimensions_for(lmax); \ } DPA4C_FOR_EACH_CHANNEL(DPA4C_DIMENSIONS) #undef DPA4C_DIMENSIONS @@ -107,6 +109,10 @@ struct Payload { torch::Tensor coupling_value; torch::Tensor output_mean; torch::Tensor output_inv_std; + // Native spin inputs. Empty together when the descriptor is spin free. + torch::Tensor spin; + torch::Tensor spin_pair; + torch::Tensor spin_type; bool canonical; int64_t lmax; double table_stride; @@ -229,7 +235,44 @@ Arguments build_arguments(const Payload& payload, "dpa4c_graph_compress: destination_row_ptr must have N + 1 " "entries"); + // === Native spin === + // The three inputs are present together or not at all. ``spin`` spans the + // absolute node axis because neighbour lookups address it with source + // indices, and ``spin_type`` packs the four per-type scalars a node reads. + const bool has_spin = payload.spin.numel() != 0; + if (has_spin) { + for (const torch::Tensor* tensor : + {&payload.spin, &payload.spin_pair, &payload.spin_type}) { + TORCH_CHECK(tensor->is_cuda() && tensor->device() == device && + tensor->is_contiguous() && + tensor->scalar_type() == torch::kFloat32, + "dpa4c_graph_compress: spin inputs must be contiguous fp32 " + "CUDA tensors on the device of edge_vec"); + } + TORCH_CHECK(payload.spin.dim() == 2 && payload.spin.size(1) == 3 && + payload.spin.size(0) == payload.atype.size(0), + "dpa4c_graph_compress: spin must have shape (N_all, 3)"); + TORCH_CHECK( + payload.spin_pair.sizes() == + torch::IntArrayRef({static_cast(type_count) * type_count, + widths.spin_channels, 2}), + "dpa4c_graph_compress: invalid ordered spin cache shape"); + TORCH_CHECK(payload.spin_type.sizes() == + torch::IntArrayRef({static_cast(type_count), 4}), + "dpa4c_graph_compress: invalid per-type spin table shape"); + } else { + TORCH_CHECK( + payload.spin_pair.numel() == 0 && payload.spin_type.numel() == 0, + "dpa4c_graph_compress: spin tables require a spin input"); + } + Arguments arguments; + arguments.has_spin = has_spin; + arguments.spin = has_spin ? payload.spin.data_ptr() : nullptr; + arguments.spin_pair = + has_spin ? payload.spin_pair.data_ptr() : nullptr; + arguments.spin_type = + has_spin ? payload.spin_type.data_ptr() : nullptr; arguments.node_count = node_count; arguments.edge_count = edge_vec.size(0); arguments.lmax = static_cast(payload.lmax); @@ -289,6 +332,9 @@ std::tuple dpa4c_graph_compress( torch::Tensor coupling_value, torch::Tensor output_mean, torch::Tensor output_inv_std, + torch::Tensor spin, + torch::Tensor spin_pair, + torch::Tensor spin_type, bool canonical, int64_t lmax, double table_stride, @@ -296,13 +342,30 @@ std::tuple dpa4c_graph_compress( double rcut, double eps, double degree_floor) { - const Payload payload{edge_index, edge_mask, destination_order, - destination_row_ptr, atype, table, - pair_film, pair_mixing, type_embedding, - readout_matrices, coupling_meta, coupling_entry, - coupling_value, output_mean, output_inv_std, - canonical, lmax, table_stride, - table_max, rcut, eps, + const Payload payload{edge_index, + edge_mask, + destination_order, + destination_row_ptr, + atype, + table, + pair_film, + pair_mixing, + type_embedding, + readout_matrices, + coupling_meta, + coupling_entry, + coupling_value, + output_mean, + output_inv_std, + spin, + spin_pair, + spin_type, + canonical, + lmax, + table_stride, + table_max, + rcut, + eps, degree_floor}; const long node_count = destination_row_ptr.numel() - 1; // The destination row pointer defines the node axis; the type table must @@ -312,7 +375,7 @@ std::tuple dpa4c_graph_compress( "different node counts"); const int channels = static_cast(type_embedding.size(1)); const Dimensions widths = - profile_dimensions(channels, static_cast(lmax)); + profile_dimensions(channels, static_cast(lmax), spin.numel() != 0); TORCH_CHECK(edge_vec.is_cuda(), "dpa4c_graph_compress: edge_vec must be a CUDA tensor"); const c10::cuda::CUDAGuard device_guard(edge_vec.device()); @@ -331,41 +394,62 @@ std::tuple dpa4c_graph_compress( return {descriptor, state}; } -torch::Tensor dpa4c_graph_compress_backward_impl( - torch::Tensor descriptor_gradient, - torch::Tensor state, - torch::Tensor edge_vec, - torch::Tensor edge_index, - torch::Tensor edge_mask, - torch::Tensor destination_order, - torch::Tensor destination_row_ptr, - torch::Tensor atype, - torch::Tensor table, - torch::Tensor pair_film, - torch::Tensor pair_mixing, - torch::Tensor type_embedding, - torch::Tensor readout_matrices, - torch::Tensor coupling_meta, - torch::Tensor coupling_entry, - torch::Tensor coupling_value, - torch::Tensor output_mean, - torch::Tensor output_inv_std, - bool canonical, - int64_t lmax, - double table_stride, - double table_max, - double rcut, - double eps, - double degree_floor, - bool reuse_state) { - const Payload payload{edge_index, edge_mask, destination_order, - destination_row_ptr, atype, table, - pair_film, pair_mixing, type_embedding, - readout_matrices, coupling_meta, coupling_entry, - coupling_value, output_mean, output_inv_std, - canonical, lmax, table_stride, - table_max, rcut, eps, +std::tuple +dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor destination_order, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + torch::Tensor spin, + torch::Tensor spin_pair, + torch::Tensor spin_type, + bool canonical, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor, + bool reuse_state) { + const Payload payload{edge_index, + edge_mask, + destination_order, + destination_row_ptr, + atype, + table, + pair_film, + pair_mixing, + type_embedding, + readout_matrices, + coupling_meta, + coupling_entry, + coupling_value, + output_mean, + output_inv_std, + spin, + spin_pair, + spin_type, + canonical, + lmax, + table_stride, + table_max, + rcut, + eps, degree_floor}; + const bool has_spin = spin.numel() != 0; const long node_count = destination_row_ptr.numel() - 1; // The destination row pointer defines the node axis; the type table must // describe exactly that axis, or the two disagree on how many nodes exist. @@ -375,7 +459,7 @@ torch::Tensor dpa4c_graph_compress_backward_impl( "different node counts"); const int channels = static_cast(type_embedding.size(1)); const Dimensions widths = - profile_dimensions(channels, static_cast(lmax)); + profile_dimensions(channels, static_cast(lmax), has_spin); TORCH_CHECK(edge_vec.is_cuda(), "dpa4c_graph_compress_backward: edge_vec must be a CUDA tensor"); const c10::cuda::CUDAGuard device_guard(edge_vec.device()); @@ -390,8 +474,16 @@ torch::Tensor dpa4c_graph_compress_backward_impl( descriptor_gradient.device() == edge_vec.device() && descriptor_gradient.numel() == node_count * widths.output_width, "dpa4c_graph_compress_backward: invalid descriptor gradient"); + auto float_options = edge_vec.options().dtype(torch::kFloat32); + // Every absent output receives its own allocation. The schema declares three + // unannotated results, so returning one empty tensor in two slots would + // introduce an alias the schema does not describe, which is undefined under + // functionalization for all three inputs rather than only for spin. + const auto absent = [&float_options] { + return torch::empty({0}, float_options); + }; if (node_count == 0) { - return torch::zeros_like(edge_vec); + return {torch::zeros_like(edge_vec), absent(), absent()}; } auto descriptor_gradient_float = descriptor_gradient.to(torch::kFloat32).contiguous(); @@ -400,47 +492,64 @@ torch::Tensor dpa4c_graph_compress_backward_impl( // The moment cotangent has exactly the layout of the saved state, so an // inference caller that no longer needs the state can reuse its storage. auto moment_gradient = reuse_state ? state : torch::empty_like(state); + // The on-site magnetic gradient closes in the node kernel; the neighbour + // part is emitted per edge and reduced onto source nodes by the shared edge + // assembly, which already walks the source CSR for the conservative force. + auto spin_gradient = + has_spin ? torch::empty({node_count, 3}, float_options) : absent(); + auto edge_spin_gradient = + has_spin ? torch::empty_like(edge_vec_float) : absent(); Arguments arguments = build_arguments(payload, edge_vec_float, channels, widths); arguments.descriptor_gradient = descriptor_gradient_float.data_ptr(); arguments.state = state.data_ptr(); arguments.moment_gradient = moment_gradient.data_ptr(); arguments.edge_gradient = edge_gradient.data_ptr(); + if (has_spin) { + arguments.spin_gradient = spin_gradient.data_ptr(); + arguments.edge_spin_gradient = edge_spin_gradient.data_ptr(); + } dispatch(channels, true, arguments, at::cuda::getCurrentCUDAStream()); - return edge_gradient.to(edge_vec.scalar_type()); + return {edge_gradient.to(edge_vec.scalar_type()), spin_gradient, + edge_spin_gradient}; } -torch::Tensor dpa4c_graph_compress_backward(torch::Tensor descriptor_gradient, - torch::Tensor state, - torch::Tensor edge_vec, - torch::Tensor edge_index, - torch::Tensor edge_mask, - torch::Tensor destination_order, - torch::Tensor destination_row_ptr, - torch::Tensor atype, - torch::Tensor table, - torch::Tensor pair_film, - torch::Tensor pair_mixing, - torch::Tensor type_embedding, - torch::Tensor readout_matrices, - torch::Tensor coupling_meta, - torch::Tensor coupling_entry, - torch::Tensor coupling_value, - torch::Tensor output_mean, - torch::Tensor output_inv_std, - bool canonical, - int64_t lmax, - double table_stride, - double table_max, - double rcut, - double eps, - double degree_floor) { +std::tuple +dpa4c_graph_compress_backward(torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor destination_order, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + torch::Tensor spin, + torch::Tensor spin_pair, + torch::Tensor spin_type, + bool canonical, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor) { return dpa4c_graph_compress_backward_impl( descriptor_gradient, state, edge_vec, edge_index, edge_mask, destination_order, destination_row_ptr, atype, table, pair_film, pair_mixing, type_embedding, readout_matrices, coupling_meta, - coupling_entry, coupling_value, output_mean, output_inv_std, canonical, - lmax, table_stride, table_max, rcut, eps, degree_floor, false); + coupling_entry, coupling_value, output_mean, output_inv_std, spin, + spin_pair, spin_type, canonical, lmax, table_stride, table_max, rcut, eps, + degree_floor, false); } std::tuple dpa4c_canonical_compress( @@ -458,6 +567,9 @@ std::tuple dpa4c_canonical_compress( torch::Tensor coupling_value, torch::Tensor output_mean, torch::Tensor output_inv_std, + torch::Tensor spin, + torch::Tensor spin_pair, + torch::Tensor spin_type, int64_t lmax, double table_stride, double table_max, @@ -469,38 +581,41 @@ std::tuple dpa4c_canonical_compress( "edge axis"); auto edge_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); auto destination_order = torch::empty({0}, source.options()); - return dpa4c_graph_compress(edge_vec, source, edge_mask, destination_order, - destination_row_ptr, atype, table, pair_film, - pair_mixing, type_embedding, readout_matrices, - coupling_meta, coupling_entry, coupling_value, - output_mean, output_inv_std, true, lmax, - table_stride, table_max, rcut, eps, degree_floor); + return dpa4c_graph_compress( + edge_vec, source, edge_mask, destination_order, destination_row_ptr, + atype, table, pair_film, pair_mixing, type_embedding, readout_matrices, + coupling_meta, coupling_entry, coupling_value, output_mean, + output_inv_std, spin, spin_pair, spin_type, true, lmax, table_stride, + table_max, rcut, eps, degree_floor); } -torch::Tensor dpa4c_canonical_compress_backward_common( - torch::Tensor descriptor_gradient, - torch::Tensor state, - torch::Tensor edge_vec, - torch::Tensor source, - torch::Tensor destination_row_ptr, - torch::Tensor atype, - torch::Tensor table, - torch::Tensor pair_film, - torch::Tensor pair_mixing, - torch::Tensor type_embedding, - torch::Tensor readout_matrices, - torch::Tensor coupling_meta, - torch::Tensor coupling_entry, - torch::Tensor coupling_value, - torch::Tensor output_mean, - torch::Tensor output_inv_std, - int64_t lmax, - double table_stride, - double table_max, - double rcut, - double eps, - double degree_floor, - bool reuse_state) { +std::tuple +dpa4c_canonical_compress_backward_common(torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + torch::Tensor spin, + torch::Tensor spin_pair, + torch::Tensor spin_type, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor, + bool reuse_state) { TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), "dpa4c_canonical_compress_backward: source and edge_vec must " "share the edge axis"); @@ -510,70 +625,77 @@ torch::Tensor dpa4c_canonical_compress_backward_common( descriptor_gradient, state, edge_vec, source, edge_mask, destination_order, destination_row_ptr, atype, table, pair_film, pair_mixing, type_embedding, readout_matrices, coupling_meta, - coupling_entry, coupling_value, output_mean, output_inv_std, true, lmax, - table_stride, table_max, rcut, eps, degree_floor, reuse_state); + coupling_entry, coupling_value, output_mean, output_inv_std, spin, + spin_pair, spin_type, true, lmax, table_stride, table_max, rcut, eps, + degree_floor, reuse_state); } -torch::Tensor dpa4c_canonical_compress_backward( - torch::Tensor descriptor_gradient, - torch::Tensor state, - torch::Tensor edge_vec, - torch::Tensor source, - torch::Tensor destination_row_ptr, - torch::Tensor atype, - torch::Tensor table, - torch::Tensor pair_film, - torch::Tensor pair_mixing, - torch::Tensor type_embedding, - torch::Tensor readout_matrices, - torch::Tensor coupling_meta, - torch::Tensor coupling_entry, - torch::Tensor coupling_value, - torch::Tensor output_mean, - torch::Tensor output_inv_std, - int64_t lmax, - double table_stride, - double table_max, - double rcut, - double eps, - double degree_floor) { +std::tuple +dpa4c_canonical_compress_backward(torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + torch::Tensor spin, + torch::Tensor spin_pair, + torch::Tensor spin_type, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor) { return dpa4c_canonical_compress_backward_common( descriptor_gradient, state, edge_vec, source, destination_row_ptr, atype, table, pair_film, pair_mixing, type_embedding, readout_matrices, coupling_meta, coupling_entry, coupling_value, output_mean, - output_inv_std, lmax, table_stride, table_max, rcut, eps, degree_floor, - false); + output_inv_std, spin, spin_pair, spin_type, lmax, table_stride, table_max, + rcut, eps, degree_floor, false); } -torch::Tensor dpa4c_canonical_compress_backward_inplace( - torch::Tensor descriptor_gradient, - torch::Tensor state, - torch::Tensor edge_vec, - torch::Tensor source, - torch::Tensor destination_row_ptr, - torch::Tensor atype, - torch::Tensor table, - torch::Tensor pair_film, - torch::Tensor pair_mixing, - torch::Tensor type_embedding, - torch::Tensor readout_matrices, - torch::Tensor coupling_meta, - torch::Tensor coupling_entry, - torch::Tensor coupling_value, - torch::Tensor output_mean, - torch::Tensor output_inv_std, - int64_t lmax, - double table_stride, - double table_max, - double rcut, - double eps, - double degree_floor) { +std::tuple +dpa4c_canonical_compress_backward_inplace(torch::Tensor descriptor_gradient, + torch::Tensor state, + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + torch::Tensor spin, + torch::Tensor spin_pair, + torch::Tensor spin_type, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor) { return dpa4c_canonical_compress_backward_common( descriptor_gradient, state, edge_vec, source, destination_row_ptr, atype, table, pair_film, pair_mixing, type_embedding, readout_matrices, coupling_meta, coupling_entry, coupling_value, output_mean, - output_inv_std, lmax, table_stride, table_max, rcut, eps, degree_floor, - true); + output_inv_std, spin, spin_pair, spin_type, lmax, table_stride, table_max, + rcut, eps, degree_floor, true); } // Energy and edge cotangent of one compressed inference step, evaluated over @@ -589,7 +711,7 @@ torch::Tensor dpa4c_canonical_compress_backward_inplace( // // The loop lives here rather than in Python because its trip count follows a // dynamic node count, which export cannot trace. -std::tuple +std::tuple dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, torch::Tensor source, torch::Tensor destination_row_ptr, @@ -604,6 +726,9 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, torch::Tensor coupling_value, torch::Tensor output_mean, torch::Tensor output_inv_std, + torch::Tensor spin, + torch::Tensor spin_pair, + torch::Tensor spin_type, int64_t lmax, double table_stride, double table_max, @@ -631,8 +756,9 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, "dpa4c_canonical_compress_energy_gradient: atype and " "destination_row_ptr describe different node counts"); const int channels = static_cast(type_embedding.size(1)); + const bool has_spin = spin.numel() != 0; const Dimensions widths = - profile_dimensions(channels, static_cast(lmax)); + profile_dimensions(channels, static_cast(lmax), has_spin); auto f32 = edge_vec.options().dtype(torch::kFloat32); auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous(); auto energy = @@ -641,8 +767,21 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, // would clear it never runs. auto edge_gradient = node_count == 0 ? torch::zeros_like(edge_vec_float) : torch::empty_like(edge_vec_float); + // Every absent output receives its own allocation, so that no two of the + // four unannotated results share storage. + const auto absent = [&f32] { return torch::empty({0}, f32); }; + // The on-site magnetic gradient is node local, so a run writes only its own + // rows. The neighbour part belongs to source nodes that other runs own, so + // it is materialized over the whole edge axis and reduced once afterwards, + // exactly like the conservative edge cotangent. + auto spin_gradient = has_spin ? torch::empty({node_count, 3}, f32) : absent(); + auto edge_spin_gradient = + has_spin ? (node_count == 0 ? torch::zeros_like(edge_vec_float) + : torch::empty_like(edge_vec_float)) + : absent(); if (node_count == 0) { - return {energy, edge_gradient.to(edge_vec.scalar_type())}; + return {energy, edge_gradient.to(edge_vec.scalar_type()), spin_gradient, + edge_spin_gradient}; } auto seed_c = seed.contiguous(); TORCH_CHECK( @@ -672,29 +811,19 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, for (long begin = 0; begin < node_count; begin += run) { const long count = std::min(run, node_count - begin); const Payload payload{ - source, - empty_mask, - empty_index, - destination_row_ptr.slice(0, begin, begin + count + 1), - atype, - table, - pair_film, - pair_mixing, - type_embedding, - readout_matrices, - coupling_meta, - coupling_entry, - coupling_value, - output_mean, - output_inv_std, - true, - lmax, - table_stride, - table_max, - rcut, - eps, - degree_floor, - begin}; + source, empty_mask, + empty_index, destination_row_ptr.slice(0, begin, begin + count + 1), + atype, table, + pair_film, pair_mixing, + type_embedding, readout_matrices, + coupling_meta, coupling_entry, + coupling_value, output_mean, + output_inv_std, spin, + spin_pair, spin_type, + true, lmax, + table_stride, table_max, + rcut, eps, + degree_floor, begin}; Arguments arguments = build_arguments(payload, edge_vec_float, channels, widths); arguments.descriptor = descriptor.data_ptr(); @@ -716,12 +845,18 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, arguments.state = state.data_ptr(); arguments.moment_gradient = state.data_ptr(); arguments.edge_gradient = edge_gradient.data_ptr(); + if (has_spin) { + // Node-indexed like the energy, so the run addresses its own slice. + arguments.spin_gradient = spin_gradient.data_ptr() + begin * 3; + arguments.edge_spin_gradient = edge_spin_gradient.data_ptr(); + } // Only the final run reaches the reserved edge slots; its row pointer ends // at the last physical edge, which is exactly where the padding begins. arguments.clear_padding = begin + count == node_count; dispatch(channels, true, arguments, stream); } - return {energy, edge_gradient.to(edge_vec.scalar_type())}; + return {energy, edge_gradient.to(edge_vec.scalar_type()), spin_gradient, + edge_spin_gradient}; } TORCH_LIBRARY_FRAGMENT(deepmd, library) { @@ -732,6 +867,7 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, " "bool canonical, int lmax, float table_stride, float table_max, " "float rcut, float eps, float degree_floor) " "-> (Tensor descriptor, Tensor state)"); @@ -743,8 +879,11 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor table, Tensor pair_film, Tensor pair_mixing, " "Tensor type_embedding, Tensor readout_matrices, Tensor coupling_meta, " "Tensor coupling_entry, Tensor coupling_value, Tensor output_mean, " - "Tensor output_inv_std, bool canonical, int lmax, float table_stride, " - "float table_max, float rcut, float eps, float degree_floor) -> Tensor"); + "Tensor output_inv_std, Tensor spin, Tensor spin_pair, " + "Tensor spin_type, bool canonical, int lmax, float table_stride, " + "float table_max, float rcut, float eps, float degree_floor) " + "-> (Tensor edge_gradient, Tensor spin_gradient, " + "Tensor edge_spin_gradient)"); library.impl("dpa4c_graph_compress_backward", torch::kCUDA, &dpa4c_graph_compress_backward); library.def( @@ -753,6 +892,7 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, " "int lmax, float table_stride, float table_max, float rcut, float eps, " "float degree_floor) -> (Tensor descriptor, Tensor state)"); library.impl("dpa4c_canonical_compress", torch::kCUDA, @@ -764,8 +904,11 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, " "int lmax, float table_stride, float table_max, float rcut, float eps, " - "float degree_floor) -> Tensor"); + "float degree_floor) " + "-> (Tensor edge_gradient, Tensor spin_gradient, " + "Tensor edge_spin_gradient)"); library.impl("dpa4c_canonical_compress_backward", torch::kCUDA, &dpa4c_canonical_compress_backward); library.def( @@ -775,8 +918,11 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, " "int lmax, float table_stride, float table_max, float rcut, float eps, " - "float degree_floor) -> Tensor"); + "float degree_floor) " + "-> (Tensor edge_gradient, Tensor spin_gradient, " + "Tensor edge_spin_gradient)"); library.impl("dpa4c_canonical_compress_backward_inplace", torch::kCUDA, &dpa4c_canonical_compress_backward_inplace); library.def( @@ -785,10 +931,13 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, " "int lmax, float table_stride, float table_max, float rcut, float eps, " "float degree_floor, Tensor[] ws, Tensor[] bs, int[] resnets, " "Tensor w_head, Tensor b_head, Tensor bias_atom_e, int act, " - "Tensor seed, int tile) -> (Tensor energy, Tensor edge_gradient)"); + "Tensor seed, int tile) " + "-> (Tensor energy, Tensor edge_gradient, Tensor spin_gradient, " + "Tensor edge_spin_gradient)"); library.impl("dpa4c_canonical_compress_energy_gradient", torch::kCUDA, &dpa4c_canonical_compress_energy_gradient); } diff --git a/source/op/pt/dpa4c_graph_compress.cuh b/source/op/pt/dpa4c_graph_compress.cuh index 7daa085a27..002fb5bd29 100644 --- a/source/op/pt/dpa4c_graph_compress.cuh +++ b/source/op/pt/dpa4c_graph_compress.cuh @@ -144,7 +144,8 @@ __device__ __forceinline__ float2 evaluate_table_with_derivative( const float4 low = __ldg(row.quartet + channel); const float2 high = __ldg(row.pair + channel); const float value = - low.x + (low.y + (low.z + (low.w + (high.x + high.y * x) * x) * x) * x) * x; + low.x + + (low.y + (low.z + (low.w + (high.x + high.y * x) * x) * x) * x) * x; const float derivative = low.y + (2.0f * low.z + (3.0f * low.w + (4.0f * high.x + 5.0f * high.y * x) * x) * x) * @@ -249,6 +250,8 @@ struct EdgeGeometry { float inverse_radius; float envelope; int source_type; + // Retained for the native spin branch, which gathers the source moment. + long source; }; template @@ -260,6 +263,7 @@ __device__ __forceinline__ EdgeGeometry load_geometry(long edge, const long* atype) { EdgeGeometry geometry; const long source = static_cast(edge_index[edge]); + geometry.source = source; geometry.source_type = static_cast(atype[source]); const float x = edge_vec[edge * 3 + 0]; const float y = edge_vec[edge * 3 + 1]; @@ -465,6 +469,93 @@ __device__ __forceinline__ void high_basis_gradient( du[2] = fmaf(weight, dz, du[2]); } +// === Native spin === + +/// Per-type scalars a spin-conditioned node reads, in one 128-bit load. +/// +/// ``scale`` is the spin gate divided by the per-type reference magnitude and +/// conditions the moment itself; ``gate`` is the bare zero-or-one flag, which +/// the magnetic-coordination family reads because that family counts +/// neighbours that carry a moment rather than the moments themselves. +struct SpinTypeWeights { + float scale; + float gate; + float vector; + float quadrupole; +}; + +__device__ __forceinline__ SpinTypeWeights load_spin_type(const float* table, + int type) { + const float4 packed = __ldg(reinterpret_cast(table) + type); + return {packed.x, packed.y, packed.z, packed.w}; +} + +/// Conditioned moment of one node, exactly zero for a non-magnetic type. +__device__ __forceinline__ void load_conditioned_spin(const float* spin, + float scale, + long node, + float (&output)[3]) { +#pragma unroll + for (int component = 0; component < 3; ++component) { + output[component] = scale * spin[node * 3 + component]; + } +} + +/// Gradient of ``B_2(s) . z`` with respect to ``s``, scaled by ``factor``. +/// +/// Writing ``Z = STF(z)``, the contraction equals ``sqrt(3/2) s^T Z s`` +/// because ``Z`` is symmetric and traceless, so the gradient is +/// ``sqrt(6) Z s``. This is the only route by which a degree-two spin +/// cotangent reaches the magnetic force. +/// +/// The six distinct entries of ``Z`` are formed as scalars rather than +/// through a ``Matrix3`` temporary: the caller sits in the innermost edge +/// loop of a kernel pinned at 64 registers, where a nine-element array is +/// enough to spill the widest profiles. ``packed`` may address shared memory. +__device__ __forceinline__ void spin_quadrupole_vjp(const float* packed, + float factor, + const float (&spin)[3], + float (&output)[3]) { + constexpr float kSqrtSix = 2.4494897427831780982f; + const float offdiagonal = factor * kSqrtSix * kInvSqrtTwo; + const float xy = offdiagonal * packed[0]; + const float yz = offdiagonal * packed[1]; + const float xz = offdiagonal * packed[3]; + const float trace = factor * kSqrtSix * kInvSqrtSix * packed[2]; + const float split = offdiagonal * packed[4]; + const float xx = split - trace; + const float yy = -split - trace; + const float zz = 2.0f * trace; + output[0] = fmaf(xx, spin[0], fmaf(xy, spin[1], xz * spin[2])); + output[1] = fmaf(xy, spin[0], fmaf(yy, spin[1], yz * spin[2])); + output[2] = fmaf(xz, spin[0], fmaf(yz, spin[1], zz * spin[2])); +} + +/// One packed degree-two component of a Cartesian vector. +/// +/// This is the same real Cartesian harmonic block the geometry uses, applied +/// to the conditioned magnetic moment. The block is a homogeneous quadratic +/// polynomial, so it is smooth at a vanishing moment and no magnitude ever +/// enters through a square root. Components are evaluated one at a time so a +/// lane that owns a single component holds no other. +__device__ __forceinline__ float spin_quadrupole_component(float x, + float y, + float z, + int component) { + switch (component) { + case 0: + return kSqrtThree * x * y; + case 1: + return kSqrtThree * y * z; + case 2: + return 0.5f * (3.0f * z * z - (x * x + y * y + z * z)); + case 3: + return kSqrtThree * x * z; + default: + return 0.5f * kSqrtThree * (x * x - y * y); + } +} + // === Symmetric traceless degree-two algebra === struct Matrix3 { @@ -579,7 +670,8 @@ __device__ __forceinline__ float readout_weight(const float* matrices, int matrix, int row, int column) { - using P = Profile; + // Geometric widths only: the readout table has no spin block. + using P = Profile; return __ldg(matrices + (static_cast(matrix) * P::C1 + row) * P::C1 + column); } @@ -593,7 +685,8 @@ __device__ __forceinline__ float probe_value(const float* probes, int degree, int component, int rank_index) { - using P = Profile; + // Geometric widths only: the probe layout has no spin block. + using P = Profile; if (degree == 1) { return probes[component * P::K1 + rank_index]; } @@ -606,21 +699,30 @@ __device__ __forceinline__ float probe_value(const float* probes, // The Cartesian basis VJP maps angular cotangents to a coordinate gradient. // Applying it per lane reduces three Cartesian components instead of the full // set of angular components across the edge group. +// +// Two kinds of angular cotangent arrive here. `d_basis` holds the cotangents +// of the degree-zero through degree-two harmonic components, whose chain rule +// to the direction is applied below in closed form. `direction_gradient` holds +// the direction cotangent of every family that evaluates its own chain rule +// beforehand -- the single-channel high degrees and the bond-projected spin +// channels -- and therefore enters as a plain sum. Both are then carried +// through the same transverse projection, which is what makes one edge produce +// one coordinate gradient rather than one per family. __device__ __forceinline__ void basis_vjp(const EdgeGeometry& geometry, const float (&d_basis)[9], - const float (&high_du)[3], + const float (&direction_gradient)[3], float radial_gradient, float (&output)[3]) { const float dux = - high_du[0] + d_basis[1] + + direction_gradient[0] + d_basis[1] + kSqrtThree * (d_basis[4] * geometry.uy + d_basis[7] * geometry.uz + d_basis[8] * geometry.ux); const float duy = - high_du[1] + d_basis[2] + + direction_gradient[1] + d_basis[2] + kSqrtThree * (d_basis[4] * geometry.ux + d_basis[5] * geometry.uz - d_basis[8] * geometry.uy); const float duz = - high_du[2] + d_basis[3] + + direction_gradient[2] + d_basis[3] + kSqrtThree * (d_basis[5] * geometry.uy + d_basis[7] * geometry.ux) + 3.0f * d_basis[6] * geometry.uz; const float dot = geometry.ux * dux + geometry.uy * duy + geometry.uz * duz; diff --git a/source/op/pt/dpa4c_graph_compress_kernel.cuh b/source/op/pt/dpa4c_graph_compress_kernel.cuh index c315be1af5..b4a62b6c1d 100644 --- a/source/op/pt/dpa4c_graph_compress_kernel.cuh +++ b/source/op/pt/dpa4c_graph_compress_kernel.cuh @@ -20,22 +20,83 @@ namespace deepmd_dpa4c { // and it keeps the concurrent groups of one warp on distinct banks. constexpr int kModeStride = kMaxRadialModes + 4; +// === Native spin block accessors === +// +// Channel zero of each spin block is the node-local on-site channel, which is +// stored outside the reduced region because it carries no neighborhood +// normalizer. Reading both through one accessor keeps the Gram loops free of +// that distinction. + +/// Flat moment coordinate of one entry of the joint degree-one spin block. +/// +/// The block holds the on-site moment at channel zero, the ``Cs`` isotropic +/// neighbour channels, and the ``Cs`` bond-projected neighbour channels, in +/// that order. The two neighbour families share the grading of the block, so +/// the Gram of the whole block is what emits every admissible degree-one spin +/// invariant; resolving all three regions through one accessor keeps the Gram +/// and its VJP free of the distinction between them. +template +__device__ __forceinline__ int spin_vector_offset(int component, int channel) { + using P = Profile; + if (channel == 0) { + return P::SpinOnsiteVector + component; + } + const int neighbor = channel - 1; + return neighbor < P::Cs + ? P::SpinVector + component * P::Cs + neighbor + : P::SpinBond + component * P::Cs + (neighbor - P::Cs); +} + +template +__device__ __forceinline__ float spin_vector_value(const float* moments, + int component, + int channel) { + return moments[spin_vector_offset(component, channel)]; +} + +template +__device__ __forceinline__ float spin_tensor_value(const float* moments, + int component, + int channel) { + using P = Profile; + return channel == 0 ? moments[P::SpinOnsiteTensor + component] + : moments[P::SpinTensor + component]; +} + +/// Decode one retained entry of the quadrupole Gram, whose on-site self-term +/// is omitted because ``|B_2(s)|^2 = |s|^4`` makes it a per-type constant +/// times the vector block's own on-site self-term. +__device__ __forceinline__ void decode_spin_tensor_pair(int entry, + int& row, + int& column) { + row = entry; + column = 1; +} + // === Forward === template -__global__ __launch_bounds__(Profile::Threads, - 32) void forward_kernel(Arguments args) { - using P = Profile; +__global__ __launch_bounds__( + Profile::Threads, + Profile:: + MinBlocksPerSm) void forward_kernel(Arguments args) { + using P = Profile; constexpr int EdgeWidth = P::ForwardEdgeWidth; constexpr int Groups = kWarpSize / EdgeWidth; constexpr int ChannelTiles = Channels / EdgeWidth; constexpr int AngularTiles = (P::C1 + EdgeWidth - 1) / EdgeWidth; constexpr int TensorTiles = (P::C2 + EdgeWidth - 1) / EdgeWidth; constexpr int HighTiles = (P::HighCount + EdgeWidth - 1) / EdgeWidth; + // The spin families share the neighbour width with degree two, so they + // inherit its tiling. The single neighbour quadrupole distributes its five + // components across lanes, exactly as the single-channel high degrees do. + constexpr int SpinTiles = HasSpin ? TensorTiles : 0; + constexpr int SpinTensorTiles = HasSpin ? (5 + EdgeWidth - 1) / EdgeWidth : 0; const int thread = threadIdx.x; const long node = blockIdx.x; @@ -60,8 +121,17 @@ __global__ __launch_bounds__(Profile::Threads, float high[HighTiles > 0 ? HighTiles : 1] = {}; float scalar_mass = 0.0f; float angular_mass = 0.0f; + float spin_magnitude[SpinTiles > 0 ? SpinTiles : 1] = {}; + float spin_coordination[SpinTiles > 0 ? SpinTiles : 1] = {}; + float spin_vector[SpinTiles > 0 ? SpinTiles : 1][3] = {}; + float spin_bond[SpinTiles > 0 ? SpinTiles : 1][3] = {}; + float spin_tensor[SpinTensorTiles > 0 ? SpinTensorTiles : 1] = {}; - __shared__ float mode_cache[HasModes ? Groups * kModeStride : 1]; + // Explicitly aligned: the mode residual reads these rows as float4, and + // a preceding shared array of arbitrary length would otherwise leave the + // block on a four-byte boundary. + __shared__ __align__( + 16) float mode_cache[HasModes ? Groups * kModeStride : 1]; float* modes = mode_cache + (HasModes ? group * kModeStride : 0); // === Step 1. Reduce the destination row into degree-wise moments === @@ -107,7 +177,37 @@ __global__ __launch_bounds__(Profile::Threads, HasModes ? args.pair_mixing + (pair * Channels + base_channel) * radial_modes : nullptr; + + // The spin payload is hoisted out of the channel loop: every channel + // multiplies the same conditioned neighbour moment and the same projection + // of that moment onto the bond, and only the amplitude varies with the + // channel. + float neighbor_spin[3] = {0.0f, 0.0f, 0.0f}; + float neighbor_magnitude = 0.0f; + float neighbor_gate = 0.0f; + float bond_alignment = 0.0f; + const float2* spin_row = nullptr; + if constexpr (HasSpin) { + const SpinTypeWeights weights = + load_spin_type(args.spin_type, geometry.source_type); + load_conditioned_spin(args.spin, weights.scale, geometry.source, + neighbor_spin); + neighbor_magnitude = neighbor_spin[0] * neighbor_spin[0] + + neighbor_spin[1] * neighbor_spin[1] + + neighbor_spin[2] * neighbor_spin[2]; + neighbor_gate = weights.gate; + // Component of the neighbour moment along the bond. The bond-projected + // family is this scalar times the unit direction, which is `basis[1..3]`. + bond_alignment = neighbor_spin[0] * basis[1] + + neighbor_spin[1] * basis[2] + + neighbor_spin[2] * basis[3]; + spin_row = + reinterpret_cast(args.spin_pair + pair * P::Cs * 2) + + base_channel; + } + float angular_zero = 0.0f; + float spin_zero = 0.0f; #pragma unroll for (int tile = 0; tile < ChannelTiles; ++tile) { const int channel = base_channel + tile * EdgeWidth; @@ -137,12 +237,50 @@ __global__ __launch_bounds__(Profile::Threads, fmaf(angular, basis[4 + component], tensor[tile][component]); } } + if constexpr (HasSpin) { + if (channel < P::Cs) { + const float2 spin_film = __ldg(spin_row + tile * EdgeWidth); + const float weight = + fmaf(spin_film.x, radial, spin_film.y) * envelope * envelope; + spin_magnitude[tile] = + fmaf(weight, neighbor_magnitude, spin_magnitude[tile]); + spin_coordination[tile] = + fmaf(weight, neighbor_gate, spin_coordination[tile]); + const float projected = weight * bond_alignment; +#pragma unroll + for (int component = 0; component < 3; ++component) { + spin_vector[tile][component] = fmaf( + weight, neighbor_spin[component], spin_vector[tile][component]); + spin_bond[tile][component] = fmaf(projected, basis[1 + component], + spin_bond[tile][component]); + } + if (channel == 0) { + spin_zero = weight; + } + } + } if constexpr (Lmax >= 3) { if (channel == 0) { angular_zero = angular; } } } + if constexpr (HasSpin) { + // The single neighbour quadrupole reads the leading spin channel only, + // so its five components are distributed across the lanes of the edge + // group and each is evaluated on the one lane that owns it. + const float weight = __shfl_sync(mask, spin_zero, leader); +#pragma unroll + for (int component = 0; component < 5; ++component) { + if (component % EdgeWidth == base_channel) { + spin_tensor[component / EdgeWidth] = + fmaf(weight, + spin_quadrupole_component(neighbor_spin[0], neighbor_spin[1], + neighbor_spin[2], component), + spin_tensor[component / EdgeWidth]); + } + } + } if constexpr (Lmax >= 3) { // Degrees three and above read only the leading channel, so their // components are distributed across the lanes of the edge group. The @@ -194,6 +332,26 @@ __global__ __launch_bounds__(Profile::Threads, high[tile] = reduce_channel_groups(high[tile]); } } + if constexpr (HasSpin) { +#pragma unroll + for (int tile = 0; tile < SpinTiles; ++tile) { + spin_magnitude[tile] = + reduce_channel_groups(spin_magnitude[tile]); + spin_coordination[tile] = + reduce_channel_groups(spin_coordination[tile]); +#pragma unroll + for (int component = 0; component < 3; ++component) { + spin_vector[tile][component] = + reduce_channel_groups(spin_vector[tile][component]); + spin_bond[tile][component] = + reduce_channel_groups(spin_bond[tile][component]); + } + } +#pragma unroll + for (int tile = 0; tile < SpinTensorTiles; ++tile) { + spin_tensor[tile] = reduce_channel_groups(spin_tensor[tile]); + } + } } __shared__ float normalizer_shared[4]; { @@ -254,6 +412,58 @@ __global__ __launch_bounds__(Profile::Threads, } } } + if constexpr (HasSpin) { + // The reduced spin families carry the same squared envelope as every + // non-scalar geometric moment and therefore share its normalizer. +#pragma unroll + for (int tile = 0; tile < SpinTiles; ++tile) { + const int channel = base_channel + tile * EdgeWidth; + if (channel < P::Cs) { + moments[P::SpinMagnitude + channel] = + spin_magnitude[tile] * angular_norm; + moments[P::SpinCoordination + channel] = + spin_coordination[tile] * angular_norm; +#pragma unroll + for (int component = 0; component < 3; ++component) { + moments[P::SpinVector + component * P::Cs + channel] = + spin_vector[tile][component] * angular_norm; + moments[P::SpinBond + component * P::Cs + channel] = + spin_bond[tile][component] * angular_norm; + } + } + } +#pragma unroll + for (int component = 0; component < 5; ++component) { + if (component % EdgeWidth == base_channel) { + moments[P::SpinTensor + component] = + spin_tensor[component / EdgeWidth] * angular_norm; + } + } + } + } + if constexpr (HasSpin) { + // The on-site families are node local. They are written after the + // division so that an invariant pairing an on-site channel with a + // neighbour channel carries exactly one neighborhood normalizer. + if (thread == 0) { + const SpinTypeWeights weights = + load_spin_type(args.spin_type, center_type); + float center_spin[3]; + load_conditioned_spin(args.spin, weights.scale, args.node_begin + node, + center_spin); +#pragma unroll + for (int component = 0; component < 3; ++component) { + moments[P::SpinOnsiteVector + component] = + weights.vector * center_spin[component]; + } +#pragma unroll + for (int component = 0; component < 5; ++component) { + moments[P::SpinOnsiteTensor + component] = + weights.quadrupole * + spin_quadrupole_component(center_spin[0], center_spin[1], + center_spin[2], component); + } + } } __syncthreads(); @@ -485,6 +695,66 @@ __global__ __launch_bounds__(Profile::Threads, product[0] * product[0] + product[1] * product[1] + product[2] * product[2]); } + + // === Step 6. Emit the spin invariants of even spin order === + if constexpr (HasSpin) { + for (int pair = thread; pair < P::SpinGramVector; pair += P::Threads) { + int row, column; + decode_upper_pair(pair, P::SpinVectorWidth, row, column); + float value = 0.0f; +#pragma unroll + for (int component = 0; component < 3; ++component) { + value = + fmaf(spin_vector_value(moments, component, row), + spin_vector_value(moments, component, column), + value); + } + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputSpinGramVector + pair, + (row == column ? 1.0f : kSqrtTwo) * value); + } + if (thread < P::SpinGramTensor) { + int row, column; + decode_spin_tensor_pair(thread, row, column); + float value = 0.0f; +#pragma unroll + for (int component = 0; component < 5; ++component) { + value = + fmaf(spin_tensor_value(moments, component, row), + spin_tensor_value(moments, component, column), + value); + } + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputSpinGramTensor + thread, + (row == column ? 1.0f : kSqrtTwo) * value); + } + // Cross Gram against the unaligned geometric degree-two moments. Both + // factors have even spin order, and with a unit direction the on-site row + // evaluates to the single-ion anisotropy sum over neighbours. + for (int output = thread; output < P::SpinCross; output += P::Threads) { + const int probe = output / P::C2; + const int channel = output % P::C2; + float value = 0.0f; +#pragma unroll + for (int component = 0; component < 5; ++component) { + value = + fmaf(spin_tensor_value(moments, component, probe), + moments[P::TensorOffset + component * P::C2 + channel], value); + } + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputSpinCross + output, + value); + } + for (int channel = thread; channel < P::Cs; channel += P::Threads) { + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, P::OutputSpinMagnitude + channel, + moments[P::SpinMagnitude + channel]); + store_descriptor(args.descriptor, args.output_mean, args.output_inv_std, + node, P::OutputWidth, + P::OutputSpinCoordination + channel, + moments[P::SpinCoordination + channel]); + } + } } // === Node readout backward === @@ -494,14 +764,14 @@ __global__ __launch_bounds__(Profile::Threads, // Their gradients need only Q0^2, Q1^2, and the symmetrized Q0 Q1 product, so // evaluating this closed form inside the node backward avoids a second probe // projection and the associated global gradient checkpoint. -template +template __device__ __forceinline__ void add_bis222_probe_gradient( int lane, long node, const Arguments& args, const float* __restrict__ probes, float (&d_tensor)[5]) { - using P = Profile; + using P = Profile; float packed_0[5]; float packed_1[5]; #pragma unroll @@ -569,7 +839,7 @@ __device__ __forceinline__ void add_bis222_probe_gradient( // The scratch is indexed by harmonic component, which a register array cannot // address without spilling, and it is private to the lane, so the reduction // stays deterministic. -template +template __device__ __forceinline__ void accumulate_coupling_gradient( long node, const Arguments& args, @@ -578,7 +848,7 @@ __device__ __forceinline__ void accumulate_coupling_gradient( int degree, int rank_index, float* __restrict__ scratch) { - using P = Profile; + using P = Profile; for (int record = 0; record < args.coupling_count; ++record) { const int* meta = args.coupling_meta + record * 8; const int degrees[3] = {meta[0], meta[1], meta[2]}; @@ -634,10 +904,10 @@ __device__ __forceinline__ void accumulate_coupling_gradient( // Four independent lane groups share one warp. An incomplete final block // aliases inactive groups to the last valid node so every thread reaches each // block-wide barrier; stores from those groups are suppressed. -template -__global__ __launch_bounds__(Profile::Threads, +template +__global__ __launch_bounds__(Profile::Threads, 2) void node_backward_kernel(Arguments args) { - using P = Profile; + using P = Profile; constexpr int MaxComponents = 9; const int thread = threadIdx.x; const int group = thread / P::NodeWidth; @@ -793,8 +1063,8 @@ __global__ __launch_bounds__(Profile::Threads, } } if constexpr (Lmax >= 3) { - accumulate_coupling_gradient(node, args, probes, moments, - 1, lane, scratch); + accumulate_coupling_gradient( + node, args, probes, moments, 1, lane, scratch); #pragma unroll for (int component = 0; component < 3; ++component) { d_vector[component] += scratch[component]; @@ -809,8 +1079,8 @@ __global__ __launch_bounds__(Profile::Threads, if (lane < P::K2) { float d_tensor[5] = {}; - add_bis222_probe_gradient(lane, node, args, - probes + 3 * P::K1, d_tensor); + add_bis222_probe_gradient( + lane, node, args, probes + 3 * P::K1, d_tensor); for (int output = lane; output < P::Bis112; output += P::K2) { int first, second; decode_upper_pair(output / P::K2, P::K1, first, second); @@ -873,8 +1143,8 @@ __global__ __launch_bounds__(Profile::Threads, } } if constexpr (Lmax >= 3) { - accumulate_coupling_gradient(node, args, probes, moments, - 2, lane, scratch); + accumulate_coupling_gradient( + node, args, probes, moments, 2, lane, scratch); #pragma unroll for (int component = 0; component < 5; ++component) { d_tensor[component] += scratch[component]; @@ -973,6 +1243,18 @@ __global__ __launch_bounds__(Profile::Threads, } } if (lane < P::C2) { + // The spin cross Gram contracts the spin quadrupoles against the + // unaligned degree-two moments, so its cotangent joins this channel here + // rather than passing through the alignment transpose above. + float cross[HasSpin ? 2 : 1]; + if constexpr (HasSpin) { +#pragma unroll + for (int probe = 0; probe < 2; ++probe) { + cross[probe] = load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, P::OutputWidth, + P::OutputSpinCross + probe * P::C2 + lane); + } + } #pragma unroll for (int component = 0; component < 5; ++component) { float value = 0.0f; @@ -982,6 +1264,15 @@ __global__ __launch_bounds__(Profile::Threads, output, lane), value); } + if constexpr (HasSpin) { +#pragma unroll + for (int probe = 0; probe < 2; ++probe) { + value = + fmaf(cross[probe], + spin_tensor_value(moments, component, probe), + value); + } + } d_moments[P::TensorOffset + component * P::C2 + lane] = value; } } @@ -994,14 +1285,125 @@ __global__ __launch_bounds__(Profile::Threads, 2.0f * load_output_gradient(args.descriptor_gradient, args.output_inv_std, node, P::OutputWidth, P::OutputGram3 + lane); - accumulate_coupling_gradient(node, args, probes, moments, - degree, 0, scratch); + accumulate_coupling_gradient( + node, args, probes, moments, degree, 0, scratch); for (int component = 0; component < count; ++component) { d_moments[offset + component] = fmaf(gram, moments[offset + component], scratch[component]); } } } + + // === Spin readout VJP and on-site magnetic gradient === + if constexpr (HasSpin) { + float onsite_vector[3] = {0.0f, 0.0f, 0.0f}; + float onsite_tensor[5] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (int channel = lane; channel < P::SpinVectorWidth; + channel += P::NodeWidth) { + float gradient[3] = {0.0f, 0.0f, 0.0f}; + for (int other = 0; other < P::SpinVectorWidth; ++other) { + const float upstream = + (channel == other ? 2.0f : kSqrtTwo) * + load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, + P::OutputWidth, + P::OutputSpinGramVector + + gram_pair_position(channel, other, P::SpinVectorWidth)); +#pragma unroll + for (int component = 0; component < 3; ++component) { + gradient[component] = + fmaf(upstream, + spin_vector_value(moments, component, other), + gradient[component]); + } + } +#pragma unroll + for (int component = 0; component < 3; ++component) { + if (channel == 0) { + // The on-site family carries no normalizer, so its cotangent leaves + // through the magnetic gradient below and must not reach the edge + // scan or the mass VJP. + onsite_vector[component] = gradient[component]; + d_moments[P::SpinOnsiteVector + component] = 0.0f; + } else { + d_moments[spin_vector_offset(component, channel)] = + gradient[component]; + } + } + } + if (lane < 2) { + float gradient[5] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (int other = 0; other < 2; ++other) { + // The on-site self-term is not emitted, so it contributes nothing. + if (lane == 0 && other == 0) { + continue; + } + const float upstream = + (lane == other ? 2.0f : kSqrtTwo) * + load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, + P::OutputWidth, + P::OutputSpinGramTensor + (lane == 1 && other == 1 ? 1 : 0)); +#pragma unroll + for (int component = 0; component < 5; ++component) { + gradient[component] = + fmaf(upstream, + spin_tensor_value(moments, component, other), + gradient[component]); + } + } + for (int channel = 0; channel < P::C2; ++channel) { + const float upstream = load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, P::OutputWidth, + P::OutputSpinCross + lane * P::C2 + channel); +#pragma unroll + for (int component = 0; component < 5; ++component) { + gradient[component] = fmaf( + upstream, moments[P::TensorOffset + component * P::C2 + channel], + gradient[component]); + } + } + if (lane == 0) { +#pragma unroll + for (int component = 0; component < 5; ++component) { + onsite_tensor[component] = gradient[component]; + d_moments[P::SpinOnsiteTensor + component] = 0.0f; + } + } else { +#pragma unroll + for (int component = 0; component < 5; ++component) { + d_moments[P::SpinTensor + component] = gradient[component]; + } + } + } + for (int channel = lane; channel < P::Cs; channel += P::NodeWidth) { + d_moments[P::SpinMagnitude + channel] = load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, P::OutputWidth, + P::OutputSpinMagnitude + channel); + d_moments[P::SpinCoordination + channel] = load_output_gradient( + args.descriptor_gradient, args.output_inv_std, node, P::OutputWidth, + P::OutputSpinCoordination + channel); + } + // Lane zero owns both on-site cotangents, so the magnetic gradient of the + // centre closes here without a shuffle. + if (active && lane == 0) { + const int center_type = + static_cast(args.atype[args.node_begin + node]); + const SpinTypeWeights weights = + load_spin_type(args.spin_type, center_type); + float center_spin[3]; + load_conditioned_spin(args.spin, weights.scale, args.node_begin + node, + center_spin); + float quadrupole[3]; + spin_quadrupole_vjp(onsite_tensor, 1.0f, center_spin, quadrupole); +#pragma unroll + for (int component = 0; component < 3; ++component) { + args.spin_gradient[node * 3 + component] = + weights.scale * fmaf(weights.quadrupole, quadrupole[component], + weights.vector * onsite_vector[component]); + } + } + } __syncthreads(); // === Normalizer VJPs === @@ -1055,11 +1457,14 @@ __global__ __launch_bounds__(Profile::Threads, template -__global__ __launch_bounds__(Profile::Threads, - 32) void edge_backward_kernel(Arguments args) { - using P = Profile; +__global__ __launch_bounds__( + Profile::Threads, + Profile:: + MinBlocksPerSm) void edge_backward_kernel(Arguments args) { + using P = Profile; constexpr int EdgeWidth = P::BackwardEdgeWidth; constexpr int Groups = kWarpSize / EdgeWidth; constexpr int ChannelTiles = Channels / EdgeWidth; @@ -1098,6 +1503,17 @@ __global__ __launch_bounds__(Profile::Threads, scalar_gradient[channel] = __ldg(args.moment_gradient + gradient_offset + P::ScalarOffset + channel); } + // The reduced spin cotangent is a node constant that every edge rereads and + // that every lane of a group needs in full, so it is staged in shared + // memory for the same reason as the scalar cotangent above. + __shared__ float spin_gradient[HasSpin ? P::SpinEdgeWidth : 1]; + if constexpr (HasSpin) { + for (int coordinate = thread; coordinate < P::SpinEdgeWidth; + coordinate += P::Threads) { + spin_gradient[coordinate] = __ldg(args.moment_gradient + gradient_offset + + P::SpinOffset + coordinate); + } + } float d_vector[AngularTiles][3] = {}; float d_tensor[TensorTiles][5] = {}; float d_high[HighTiles > 0 ? HighTiles : 1] = {}; @@ -1135,13 +1551,25 @@ __global__ __launch_bounds__(Profile::Threads, } } - __shared__ float mode_cache[HasModes ? Groups * kModeStride : 1]; - __shared__ float mode_derivative_cache[HasModes ? Groups * kModeStride : 1]; + // Explicitly aligned: the mode residual reads these rows as float4, and + // a preceding shared array of arbitrary length would otherwise leave the + // block on a four-byte boundary. + __shared__ __align__( + 16) float mode_cache[HasModes ? Groups * kModeStride : 1]; + __shared__ __align__( + 16) float mode_derivative_cache[HasModes ? Groups * kModeStride : 1]; const int mode_offset = HasModes ? group * kModeStride : 0; float* modes = mode_cache + mode_offset; float* mode_derivatives = mode_derivative_cache + mode_offset; __syncthreads(); + // Offsets of the five spin families inside the staged cotangent row. + constexpr int SpinMagnitudeSlot = P::SpinMagnitude - P::SpinOffset; + constexpr int SpinCoordinationSlot = P::SpinCoordination - P::SpinOffset; + constexpr int SpinVectorSlot = P::SpinVector - P::SpinOffset; + constexpr int SpinBondSlot = P::SpinBond - P::SpinOffset; + constexpr int SpinTensorSlot = P::SpinTensor - P::SpinOffset; + for (long position = begin + group; position < end; position += Groups) { const long edge = edge_at_position(position, destination_order); if (args.edge_mask != nullptr && !args.edge_mask[edge]) { @@ -1149,6 +1577,11 @@ __global__ __launch_bounds__(Profile::Threads, args.edge_gradient[edge * 3 + 0] = 0.0f; args.edge_gradient[edge * 3 + 1] = 0.0f; args.edge_gradient[edge * 3 + 2] = 0.0f; + if constexpr (HasSpin) { + args.edge_spin_gradient[edge * 3 + 0] = 0.0f; + args.edge_spin_gradient[edge * 3 + 1] = 0.0f; + args.edge_spin_gradient[edge * 3 + 2] = 0.0f; + } } continue; } @@ -1164,6 +1597,11 @@ __global__ __launch_bounds__(Profile::Threads, args.edge_gradient[edge * 3 + 0] = 0.0f; args.edge_gradient[edge * 3 + 1] = 0.0f; args.edge_gradient[edge * 3 + 2] = 0.0f; + if constexpr (HasSpin) { + args.edge_spin_gradient[edge * 3 + 0] = 0.0f; + args.edge_spin_gradient[edge * 3 + 1] = 0.0f; + args.edge_spin_gradient[edge * 3 + 2] = 0.0f; + } } continue; } @@ -1210,9 +1648,43 @@ __global__ __launch_bounds__(Profile::Threads, high_angular = subwarp_sum(high_angular, mask); } + // The spin payload of this edge, hoisted like the forward one: every + // channel multiplies the same conditioned neighbour moment and the same + // projection of that moment onto the bond. + float neighbor_spin[3] = {0.0f, 0.0f, 0.0f}; + float neighbor_magnitude = 0.0f; + float neighbor_gate = 0.0f; + float neighbor_scale = 0.0f; + float bond_alignment = 0.0f; + float spin_cotangent[3] = {0.0f, 0.0f, 0.0f}; + const float2* spin_row = nullptr; + if constexpr (HasSpin) { + const SpinTypeWeights weights = + load_spin_type(args.spin_type, geometry.source_type); + load_conditioned_spin(args.spin, weights.scale, geometry.source, + neighbor_spin); + neighbor_magnitude = neighbor_spin[0] * neighbor_spin[0] + + neighbor_spin[1] * neighbor_spin[1] + + neighbor_spin[2] * neighbor_spin[2]; + neighbor_gate = weights.gate; + neighbor_scale = weights.scale; + bond_alignment = neighbor_spin[0] * basis[1] + + neighbor_spin[1] * basis[2] + + neighbor_spin[2] * basis[3]; + spin_row = + reinterpret_cast(args.spin_pair + pair * P::Cs * 2) + + base_channel; + } + float radial_gradient = 0.0f; float envelope_gradient = 0.0f; float d_basis[9] = {}; + // Direction cotangent of every family that evaluates its own angular chain + // rule: the single-channel high degrees, which fill it after the channel + // loop, and the bond-projected spin channels, which fill it inside. Both + // reach the coordinate gradient through the one transverse projection in + // `basis_vjp`, so no family is reduced separately. + float direction_gradient[3] = {0.0f, 0.0f, 0.0f}; float angular_zero = 0.0f; #pragma unroll for (int tile = 0; tile < ChannelTiles; ++tile) { @@ -1274,6 +1746,101 @@ __global__ __launch_bounds__(Profile::Threads, d_basis[4 + component]); } } + if constexpr (HasSpin) { + if (channel < P::Cs) { + const float2 spin_film = __ldg(spin_row + tile * EdgeWidth); + const float profile = fmaf(spin_film.x, radial.x, spin_film.y); + const float envelope_squared = envelope * envelope; + const float weight = profile * envelope_squared; + const float magnitude_gradient = + spin_gradient[SpinMagnitudeSlot + channel]; + // The two degree-one cotangents are read once and used by the + // amplitude, the magnetic and the angular accumulations below. + float vector_gradient[3]; + float bond_gradient[3]; +#pragma unroll + for (int component = 0; component < 3; ++component) { + vector_gradient[component] = + spin_gradient[SpinVectorSlot + component * P::Cs + channel]; + bond_gradient[component] = + spin_gradient[SpinBondSlot + component * P::Cs + channel]; + } + // Component of the bond cotangent along the bond. The bond-projected + // family is `P = w (s . u) u`, so this one dot product carries all + // three of its VJPs: `dP/dw` contracts to `(s . u) (dP . u)`, + // `dP/ds` to `w (dP . u) u`, and the direction term below reuses it + // once more. + float bond_projection = 0.0f; +#pragma unroll + for (int component = 0; component < 3; ++component) { + bond_projection = fmaf(bond_gradient[component], + basis[1 + component], bond_projection); + } + float sigma = fmaf( + magnitude_gradient, neighbor_magnitude, + spin_gradient[SpinCoordinationSlot + channel] * neighbor_gate); + sigma = fmaf(bond_alignment, bond_projection, sigma); +#pragma unroll + for (int component = 0; component < 3; ++component) { + sigma = fmaf(vector_gradient[component], neighbor_spin[component], + sigma); + } + // Magnetic cotangent of the source moment. The coordination family + // reads the type gate rather than the moment, so it contributes + // nothing here. + const float linear = 2.0f * weight * magnitude_gradient; + const float projected = weight * bond_projection; +#pragma unroll + for (int component = 0; component < 3; ++component) { + spin_cotangent[component] = + fmaf(linear, neighbor_spin[component], + fmaf(projected, basis[1 + component], + fmaf(weight, vector_gradient[component], + spin_cotangent[component]))); + } + // Angular cotangent of the bond-projected family. Differentiating + // `(s . u) u` with respect to `u` at fixed `s` gives the rank-one + // sum `s (dP . u) + (s . u) dP`, which the transverse projection in + // `basis_vjp` turns into the coordinate gradient. Every other spin + // family is independent of the direction and contributes nothing. + const float aligned_weight = weight * bond_alignment; +#pragma unroll + for (int component = 0; component < 3; ++component) { + direction_gradient[component] = + fmaf(projected, neighbor_spin[component], + fmaf(aligned_weight, bond_gradient[component], + direction_gradient[component])); + } + if (channel == 0) { + // The single neighbour quadrupole rides the leading spin channel. + // Both of its contributions are formed here rather than hoisted, + // so its five cotangent components never stay live across the + // channel loop. +#pragma unroll + for (int component = 0; component < 5; ++component) { + sigma = fmaf( + spin_gradient[SpinTensorSlot + component], + spin_quadrupole_component(neighbor_spin[0], neighbor_spin[1], + neighbor_spin[2], component), + sigma); + } + float quadrupole[3]; + spin_quadrupole_vjp(spin_gradient + SpinTensorSlot, weight, + neighbor_spin, quadrupole); +#pragma unroll + for (int component = 0; component < 3; ++component) { + spin_cotangent[component] += quadrupole[component]; + } + } + // w = chi^2 (gamma g + beta): the distance enters through the table + // and through the envelope, and `sigma` is the cotangent of that + // shared amplitude over all five families. + radial_gradient = fmaf(envelope_squared * sigma, + spin_film.x * radial.y, radial_gradient); + envelope_gradient = + fmaf(2.0f * envelope * profile, sigma, envelope_gradient); + } + } if constexpr (Lmax >= 3) { if (channel == 0) { angular_zero = angular_payload; @@ -1281,7 +1848,6 @@ __global__ __launch_bounds__(Profile::Threads, } } - float high_du[3] = {0.0f, 0.0f, 0.0f}; if constexpr (Lmax >= 3) { const float amplitude = __shfl_sync(mask, angular_zero, leader); #pragma unroll @@ -1289,7 +1855,7 @@ __global__ __launch_bounds__(Profile::Threads, if (component % EdgeWidth == base_channel) { high_basis_gradient(geometry, component, d_high[component / EdgeWidth] * amplitude, - high_du); + direction_gradient); } } } @@ -1310,7 +1876,7 @@ __global__ __launch_bounds__(Profile::Threads, // applying it per lane reduces three Cartesian components instead of the // full set of angular components across the edge group. float output[3]; - basis_vjp(geometry, d_basis, high_du, radial_gradient, output); + basis_vjp(geometry, d_basis, direction_gradient, radial_gradient, output); #pragma unroll for (int component = 0; component < 3; ++component) { output[component] = subwarp_sum(output[component], mask); @@ -1320,6 +1886,28 @@ __global__ __launch_bounds__(Profile::Threads, args.edge_gradient[edge * 3 + 1] = output[1]; args.edge_gradient[edge * 3 + 2] = output[2]; } + if constexpr (HasSpin) { + // The magnetic cotangent belongs to the source node, which this + // destination-major scan does not own. It is emitted per edge and + // reduced onto sources by the shared edge assembly, which already walks + // the source CSR for the conservative force. +#pragma unroll + for (int component = 0; component < 3; ++component) { + spin_cotangent[component] = + subwarp_sum(spin_cotangent[component], mask); + } + if (thread == leader) { + // Every contribution above differentiates the conditioned moment, so + // the store is the single point that applies the remaining chain + // factor and hands back a gradient of the raw input moment. + args.edge_spin_gradient[edge * 3 + 0] = + neighbor_scale * spin_cotangent[0]; + args.edge_spin_gradient[edge * 3 + 1] = + neighbor_scale * spin_cotangent[1]; + args.edge_spin_gradient[edge * 3 + 2] = + neighbor_scale * spin_cotangent[2]; + } + } } } @@ -1328,7 +1916,8 @@ __global__ void zero_padding_kernel(long node_count, long edge_count, const index_t* destination_order, const long* destination_row_ptr, - float* edge_gradient) { + float* edge_gradient, + float* edge_spin_gradient) { const long valid_edge_count = destination_row_ptr[node_count]; for (long position = valid_edge_count + blockIdx.x * blockDim.x + threadIdx.x; position < edge_count; @@ -1337,6 +1926,11 @@ __global__ void zero_padding_kernel(long node_count, edge_gradient[edge * 3 + 0] = 0.0f; edge_gradient[edge * 3 + 1] = 0.0f; edge_gradient[edge * 3 + 2] = 0.0f; + if (edge_spin_gradient != nullptr) { + edge_spin_gradient[edge * 3 + 0] = 0.0f; + edge_spin_gradient[edge * 3 + 1] = 0.0f; + edge_spin_gradient[edge * 3 + 2] = 0.0f; + } } } @@ -1345,12 +1939,13 @@ __global__ void zero_padding_kernel(long node_count, template struct ForwardLauncher { static void run(const Arguments& args, cudaStream_t stream) { - using P = Profile; - forward_kernel + using P = Profile; + forward_kernel <<(args.node_count), P::Threads, 0, stream>>>(args); } }; @@ -1358,16 +1953,17 @@ struct ForwardLauncher { template struct BackwardLauncher { static void run(const Arguments& args, cudaStream_t stream) { - using P = Profile; + using P = Profile; const int node_blocks = static_cast((args.node_count + P::NodeGroups - 1) / P::NodeGroups); - node_backward_kernel + node_backward_kernel <<>>(args); - edge_backward_kernel + edge_backward_kernel <<(args.node_count), P::Threads, 0, stream>>>(args); // The reserved edge slots beyond the physical count are only known on the // device, so the grid is sized from the storage bound and the surplus @@ -1382,7 +1978,8 @@ struct BackwardLauncher { <<(padding_blocks), kPaddingThreads, 0, stream>>>( args.node_count, args.edge_count, static_cast(args.destination_order), - args.destination_row_ptr, args.edge_gradient); + args.destination_row_ptr, args.edge_gradient, + HasSpin ? args.edge_spin_gradient : nullptr); } } }; @@ -1393,35 +1990,55 @@ struct BackwardLauncher { template class L> + bool HasSpin, + template class L> void dispatch_topology(const Arguments& args, cudaStream_t stream) { const bool wide = args.index_kind == IndexKind::Bits64; if (args.canonical) { if (wide) { - L::run(args, stream); + L::run(args, stream); } else { - L::run(args, stream); + L::run(args, + stream); } } else { if (wide) { - L::run(args, stream); + L::run(args, stream); } else { - L::run(args, stream); + L::run(args, + stream); } } } +// Native spin is a compile-time specialization for the same reason as the +// mode residual: a spinless descriptor must not carry the spin accumulators, +// the staged spin cotangent, or the wider moment state of one that has them. +// The neighbour spin width is derived from the degree profile, so presence is +// the whole choice and the instantiation matrix only doubles. +template class L> +void dispatch_spin(const Arguments& args, cudaStream_t stream) { + if (args.has_spin) { + dispatch_topology(args, stream); + } else { + dispatch_topology(args, stream); + } +} + // The mode residual is a compile-time specialization for the same reason as // the angular degree: a descriptor without radial modes must not carry the // vector temporaries and the shared profile cache of one that has them. template class L> + template class L> void dispatch_modes(const Arguments& args, cudaStream_t stream) { if (args.radial_modes > 0) { - dispatch_topology(args, stream); + dispatch_spin(args, stream); } else { - dispatch_topology(args, stream); + dispatch_spin(args, stream); } } @@ -1433,7 +2050,7 @@ void dispatch_modes(const Arguments& args, cudaStream_t stream) { // dispatch. The unreachable default is still checked rather than folded into // the highest degree, so that a degree outside the compiled set can only ever // fail loudly instead of running a kernel for a different model. -template class L> +template class L> void dispatch_degree(const Arguments& args, cudaStream_t stream) { switch (args.lmax) { case 2: diff --git a/source/op/pt/dpa4c_graph_compress_launch.h b/source/op/pt/dpa4c_graph_compress_launch.h index 0cb003dad9..be885ac78f 100644 --- a/source/op/pt/dpa4c_graph_compress_launch.h +++ b/source/op/pt/dpa4c_graph_compress_launch.h @@ -50,6 +50,9 @@ struct Arguments { float eps = 0.0f; float degree_floor = 0.0f; bool canonical = false; + // Whether the native spin branch is active. The neighbour spin width is + // derived from the degree profile, so presence is the whole runtime choice. + bool has_spin = false; IndexKind index_kind = IndexKind::Bits64; const void* edge_index = nullptr; @@ -71,10 +74,26 @@ struct Arguments { const float* descriptor_gradient = nullptr; const float* state = nullptr; + // === Native spin === + // Present together or not at all. ``spin`` is indexed by absolute node id, + // like ``atype``, because neighbour lookups address it with source indices. + // ``spin_type`` packs the four per-type scalars a node needs -- the fused + // gate over reference magnitude, the bare gate that the magnetic + // coordination family counts, and the two on-site weights -- so one node + // reads them in a single 128-bit load. + const float* spin = nullptr; + const float* spin_pair = nullptr; + const float* spin_type = nullptr; + float* descriptor = nullptr; float* state_out = nullptr; float* moment_gradient = nullptr; float* edge_gradient = nullptr; + // Per-node on-site magnetic gradient and the per-edge contribution to the + // source node's magnetic gradient. The latter is reduced onto source nodes + // by the shared edge force assembly, which already walks the source CSR. + float* spin_gradient = nullptr; + float* edge_spin_gradient = nullptr; }; // === Compile-time descriptor profile === @@ -168,6 +187,23 @@ constexpr int coupling_record_count(int lmax) { // measured optimum of that trade-off on a diamond neighborhood; forward and // backward differ because the backward carries the additional angular // cotangents. +// +// The native spin branch carries its own pair of widths, because its channel +// count is tied to the degree-two width rather than to the scalar width and +// its five families add accumulators that the geometric optimum does not +// account for. Both were measured on the same diamond neighborhood as the +// geometric pair, independently of each other, since one governs the forward +// kernel and the other the backward kernel. +// +// The measured spin optima differ from the geometric ones in exactly two +// places. At the narrowest profile the spin forward wants one step wider, +// because the eight-fold spin payload of a two-lane group tiles into more +// accumulators than the recovered edge concurrency is worth. At ``Channels == +// 64`` the spin backward wants one step narrower, because the bond-projected +// family adds an angular cotangent whose per-lane cost outweighs the geometry +// recompute that a wider group amortizes. A group wider than the scalar width +// leaves a tile empty and does not compile, which bounds the narrow profiles +// from above. template struct EdgeMap; @@ -175,29 +211,42 @@ template <> struct EdgeMap<8> { static constexpr int Forward = 2; static constexpr int Backward = 2; + static constexpr int SpinForward = 4; + static constexpr int SpinBackward = 2; }; template <> struct EdgeMap<16> { static constexpr int Forward = 4; static constexpr int Backward = 4; + static constexpr int SpinForward = 4; + static constexpr int SpinBackward = 4; }; template <> struct EdgeMap<32> { static constexpr int Forward = 8; static constexpr int Backward = 4; + static constexpr int SpinForward = 8; + static constexpr int SpinBackward = 4; }; template <> struct EdgeMap<64> { static constexpr int Forward = 8; static constexpr int Backward = 8; + static constexpr int SpinForward = 8; + static constexpr int SpinBackward = 4; }; template <> struct EdgeMap<128> { static constexpr int Forward = 16; static constexpr int Backward = 8; + static constexpr int SpinForward = 16; + static constexpr int SpinBackward = 8; }; -template +// ``HasSpin`` is deliberately without a default. It changes the moment +// layout and the descriptor width, so an instantiation that omits it would +// silently read a spin-free layout out of a spin-conditioned buffer. +template struct Profile { static constexpr int C0 = Channels; static constexpr int C1 = degree_one_width(Channels); @@ -205,8 +254,26 @@ struct Profile { static constexpr int K1 = C2; static constexpr int K2 = 2; - // Flat moment layout: degree zero, degree one, degree two, then the - // single-channel high degrees in increasing order. + // === Native spin widths === + // The neighbour spin width is derived from the degree-two width, so the + // presence of the branch is the only new compile-time degree of freedom and + // the instantiation matrix doubles rather than growing by a factor of four. + static constexpr int Cs = HasSpin ? C2 : 0; + // Reduced families: neighbour moment magnitude, magnetic coordination, the + // isotropic neighbour spin vector, the bond-projected neighbour spin vector, + // and the single neighbour spin quadrupole. + static constexpr int SpinEdgeWidth = HasSpin ? 8 * Cs + 5 : 0; + // Node-local on-site vector and quadrupole, written outside the division. + static constexpr int SpinNodeWidth = HasSpin ? 8 : 0; + static constexpr int SpinMomentWidth = SpinEdgeWidth + SpinNodeWidth; + // Channel width of the joint degree-one spin block: the on-site moment, the + // isotropic neighbour channels and the bond-projected neighbour channels. + // The two neighbour families share one grading, so one Gram over the joint + // block emits every admissible degree-one spin invariant. + static constexpr int SpinVectorWidth = HasSpin ? 1 + 2 * Cs : 0; + + // Flat moment layout: degree zero, degree one, degree two, the + // single-channel high degrees in increasing order, then the spin families. static constexpr int ScalarOffset = 0; static constexpr int VectorOffset = C0; static constexpr int TensorOffset = C0 + 3 * C1; @@ -214,7 +281,15 @@ struct Profile { static constexpr int High3 = Lmax >= 3 ? 7 : 0; static constexpr int High4 = Lmax >= 4 ? 9 : 0; static constexpr int HighCount = High3 + High4; - static constexpr int MomentWidth = HighOffset + HighCount; + static constexpr int SpinOffset = HighOffset + HighCount; + static constexpr int SpinMagnitude = SpinOffset; // Cs + static constexpr int SpinCoordination = SpinMagnitude + Cs; // Cs + static constexpr int SpinVector = SpinCoordination + Cs; // 3 * Cs + static constexpr int SpinBond = SpinVector + 3 * Cs; // 3 * Cs + static constexpr int SpinTensor = SpinBond + 3 * Cs; // 5 + static constexpr int SpinOnsiteVector = SpinTensor + 5; // 3 + static constexpr int SpinOnsiteTensor = SpinOnsiteVector + 3; // 5 + static constexpr int MomentWidth = SpinOffset + SpinMomentWidth; static constexpr int StateWidth = MomentWidth + 2; // Cached intermediates of the invariant readout. @@ -239,10 +314,25 @@ struct Profile { BispectrumBase + bispectrum_prefix(Lmax, K1, K2, 2, 2, 2); static constexpr int OutputQuartic = BispectrumBase + bispectrum_prefix(Lmax, K1, K2, 0, 0, 0); + // Spin invariants of even spin order. The quadrupole Gram omits its + // on-site self-term, which the identity |B_2(s)|^2 = |s|^4 makes a per-type + // constant times the square of the vector block's on-site self-term. + static constexpr int SpinGramVector = + HasSpin ? triangular(SpinVectorWidth) : 0; + static constexpr int SpinGramTensor = HasSpin ? 2 : 0; + static constexpr int SpinCross = HasSpin ? 2 * C2 : 0; + static constexpr int SpinDim = + SpinGramVector + SpinGramTensor + SpinCross + 2 * Cs; + static constexpr int OutputSpinGramVector = OutputQuartic + Quartic; + static constexpr int OutputSpinGramTensor = + OutputSpinGramVector + SpinGramVector; + static constexpr int OutputSpinCross = OutputSpinGramTensor + SpinGramTensor; + static constexpr int OutputSpinMagnitude = OutputSpinCross + SpinCross; + static constexpr int OutputSpinCoordination = OutputSpinMagnitude + Cs; // The two moment divisors close the geometric block. Normalization is // otherwise irreversible, so without them neither the readout nor the // fitting network can see the effective coordination they encode. - static constexpr int OutputDivisor = OutputQuartic + Quartic; + static constexpr int OutputDivisor = OutputQuartic + Quartic + SpinDim; static constexpr int OutputType = OutputDivisor + 2; static constexpr int OutputWidth = OutputType + C0; @@ -250,11 +340,28 @@ struct Profile { // single-channel components add one accumulator per lane and tile, but every // measured widening lost more to the reduced edge concurrency than it // recovered in register pressure. - static constexpr int ForwardEdgeWidth = EdgeMap::Forward; - static constexpr int BackwardEdgeWidth = EdgeMap::Backward; + static constexpr int ForwardEdgeWidth = + HasSpin ? EdgeMap::SpinForward : EdgeMap::Forward; + static constexpr int BackwardEdgeWidth = + HasSpin ? EdgeMap::SpinBackward : EdgeMap::Backward; static constexpr int NodeWidth = 8; static constexpr int NodeGroups = kWarpSize / NodeWidth; static constexpr int Threads = kWarpSize; + + // Resident blocks the edge kernels are compiled for. A block is one warp, so + // thirty-two of them exhaust the 65,536-register file at sixty-four + // registers per thread, and a lower target raises the per-thread budget in + // proportion at the cost of occupancy. + // + // Degree two holds its geometric working set in that budget and spills only + // what the five spin families add, so twenty-four blocks buy eighty + // registers, remove the spill outright and are worth 2.0% to 4.5% of the + // step. Degree three overflows the budget on geometry alone -- it spills + // without the spin branch and still spills at eighty and at ninety-six + // registers -- so the occupancy it would give up buys an incomplete fix and + // costs 7.4% to 9.6%. The relief is therefore taken only where it is + // complete. + static constexpr int MinBlocksPerSm = (HasSpin && Lmax == 2) ? 24 : 32; }; /// Scalar widths that own a compiled specialization. diff --git a/source/op/pt/edge_force_virial.cu b/source/op/pt/edge_force_virial.cu index 5dff4b4990..c512d4a651 100644 --- a/source/op/pt/edge_force_virial.cu +++ b/source/op/pt/edge_force_virial.cu @@ -20,6 +20,7 @@ #include #include #include +#include namespace { @@ -50,7 +51,7 @@ __global__ void build_source_order_kernel(long valid_edge_count, } } -template +template __global__ void edge_force_virial_kernel( long node_count, const scalar_t* __restrict__ edge_gradient, @@ -60,8 +61,10 @@ __global__ void edge_force_virial_kernel( const long* __restrict__ destination_row_ptr, const index_t* __restrict__ source_order, const long* __restrict__ source_row_ptr, + const scalar_t* __restrict__ edge_spin_gradient, scalar_t* __restrict__ force, - scalar_t* __restrict__ node_virial) { + scalar_t* __restrict__ node_virial, + scalar_t* __restrict__ magnetic_force) { constexpr unsigned kWarpMask = 0xffffffffu; const int lane = threadIdx.x & 31; const int warp = threadIdx.x >> 5; @@ -73,6 +76,13 @@ __global__ void edge_force_virial_kernel( scalar_t source_x = 0; scalar_t source_y = 0; scalar_t source_z = 0; + // The magnetic cotangent of an edge belongs to the node that sources it, the + // same grouping the force reduction already walks, so it rides the source + // loop and the warp fold rather than a second pass over the edge axis. The + // geometric instantiation carries neither the accumulators nor the fold. + scalar_t spin_x = 0; + scalar_t spin_y = 0; + scalar_t spin_z = 0; scalar_t virial[9] = {}; if (node < node_count) { @@ -108,6 +118,11 @@ __global__ void edge_force_virial_kernel( source_x += gx; source_y += gy; source_z += gz; + if constexpr (HasSpin) { + spin_x += edge_spin_gradient[edge * 3 + 0]; + spin_y += edge_spin_gradient[edge * 3 + 1]; + spin_z += edge_spin_gradient[edge * 3 + 2]; + } virial[0] = fma(-gx, x, virial[0]); virial[1] = fma(-gx, y, virial[1]); virial[2] = fma(-gx, z, virial[2]); @@ -128,6 +143,11 @@ __global__ void edge_force_virial_kernel( source_x += __shfl_down_sync(kWarpMask, source_x, offset); source_y += __shfl_down_sync(kWarpMask, source_y, offset); source_z += __shfl_down_sync(kWarpMask, source_z, offset); + if constexpr (HasSpin) { + spin_x += __shfl_down_sync(kWarpMask, spin_x, offset); + spin_y += __shfl_down_sync(kWarpMask, spin_y, offset); + spin_z += __shfl_down_sync(kWarpMask, spin_z, offset); + } #pragma unroll for (int component = 0; component < 9; ++component) { virial[component] += @@ -144,6 +164,11 @@ __global__ void edge_force_virial_kernel( for (int component = 0; component < 9; ++component) { output[component] = virial[component]; } + if constexpr (HasSpin) { + magnetic_force[node * 3 + 0] = spin_x; + magnetic_force[node * 3 + 1] = spin_y; + magnetic_force[node * 3 + 2] = spin_z; + } } } @@ -275,23 +300,37 @@ void launch_force_virial(long node_count, const torch::Tensor& source_order, const torch::Tensor& source_row_ptr, const torch::Tensor& frame_row_ptr, + const torch::Tensor& edge_spin_gradient, torch::Tensor& force, torch::Tensor& node_virial, + torch::Tensor& magnetic_force, torch::Tensor& virial_partial, torch::Tensor& virial, cudaStream_t stream) { const int node_blocks = static_cast((node_count + kWarpsPerBlock - 1) / kWarpsPerBlock); - edge_force_virial_kernel - <<>>( - node_count, edge_gradient.data_ptr(), - edge_vec.data_ptr(), - edge_mask.numel() ? edge_mask.data_ptr() : nullptr, - destination_order.numel() ? destination_order.data_ptr() - : nullptr, - destination_row_ptr.data_ptr(), - source_order.data_ptr(), source_row_ptr.data_ptr(), - force.data_ptr(), node_virial.data_ptr()); + const bool has_spin = edge_spin_gradient.numel() != 0; + auto assemble = [&](auto spin_tag) { + edge_force_virial_kernel + <<>>( + node_count, edge_gradient.data_ptr(), + edge_vec.data_ptr(), + edge_mask.numel() ? edge_mask.data_ptr() : nullptr, + destination_order.numel() ? destination_order.data_ptr() + : nullptr, + destination_row_ptr.data_ptr(), + source_order.data_ptr(), source_row_ptr.data_ptr(), + edge_spin_gradient.numel() ? edge_spin_gradient.data_ptr() + : nullptr, + force.data_ptr(), node_virial.data_ptr(), + magnetic_force.numel() ? magnetic_force.data_ptr() + : nullptr); + }; + if (has_spin) { + assemble(std::true_type{}); + } else { + assemble(std::false_type{}); + } FORCE_CHECK_LAUNCH("edge_force_virial node reduction"); launch_frame_segment_sum( @@ -299,17 +338,18 @@ void launch_force_virial(long node_count, virial_partial, virial.data_ptr(), stream); } -std::tuple assemble_force_virial( - long node_count, - const torch::Tensor& edge_gradient, - const torch::Tensor& edge_vec, - const torch::Tensor& edge_mask, - const torch::Tensor& destination_order, - const torch::Tensor& destination_row_ptr, - const torch::Tensor& source_order, - const torch::Tensor& source_row_ptr, - const torch::Tensor& n_node_per_frame, - bool want_atom_virial) { +std::tuple +assemble_force_virial(long node_count, + const torch::Tensor& edge_gradient, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_order, + const torch::Tensor& source_row_ptr, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& edge_spin_gradient, + bool want_atom_virial) { const long frame_count = n_node_per_frame.size(0); auto options = edge_gradient.options(); auto force = torch::empty({node_count, 3}, options); @@ -319,8 +359,10 @@ std::tuple assemble_force_virial( ? atom_virial : torch::empty({node_count, 3, 3}, options); auto virial = torch::zeros({frame_count, 3, 3}, options); + auto magnetic_force = + torch::empty({edge_spin_gradient.numel() ? node_count : 0, 3}, options); if (node_count == 0 || frame_count == 0) { - return {force, atom_virial, virial}; + return {force, atom_virial, virial, magnetic_force}; } auto frame_row_ptr = frame_row_pointer(n_node_per_frame); @@ -335,23 +377,23 @@ std::tuple assemble_force_virial( launch_force_virial( node_count, frame_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, frame_row_ptr, force, node_virial, virial_partial, - virial, stream); + source_row_ptr, frame_row_ptr, edge_spin_gradient, force, + node_virial, magnetic_force, virial_partial, virial, stream); } else if (source_order.scalar_type() == torch::kUInt32) { launch_force_virial( node_count, frame_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, frame_row_ptr, force, node_virial, virial_partial, - virial, stream); + source_row_ptr, frame_row_ptr, edge_spin_gradient, force, + node_virial, magnetic_force, virial_partial, virial, stream); } else { launch_force_virial( node_count, frame_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, frame_row_ptr, force, node_virial, virial_partial, - virial, stream); + source_row_ptr, frame_row_ptr, edge_spin_gradient, force, + node_virial, magnetic_force, virial_partial, virial, stream); } }); - return {force, atom_virial, virial}; + return {force, atom_virial, virial, magnetic_force}; } } // namespace @@ -402,18 +444,19 @@ build_graph_csr(torch::Tensor edge_index, return {destination_order, destination_row_ptr, source_order, source_row_ptr}; } -std::tuple edge_force_virial( - torch::Tensor edge_gradient, - torch::Tensor edge_vec, - torch::Tensor edge_index, - torch::Tensor edge_mask, - torch::Tensor destination_order, - torch::Tensor destination_row_ptr, - torch::Tensor source_order, - torch::Tensor source_row_ptr, - torch::Tensor n_node_per_frame, - c10::SymInt node_capacity, - bool want_atom_virial) { +std::tuple +edge_force_virial(torch::Tensor edge_gradient, + torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor destination_order, + torch::Tensor destination_row_ptr, + torch::Tensor source_order, + torch::Tensor source_row_ptr, + torch::Tensor n_node_per_frame, + torch::Tensor edge_spin_gradient, + c10::SymInt node_capacity, + bool want_atom_virial) { const long node_count = node_capacity.expect_int(); TORCH_CHECK(edge_gradient.is_cuda() && edge_vec.is_cuda() && edge_mask.is_cuda() && destination_order.is_cuda() && @@ -442,19 +485,28 @@ std::tuple edge_force_virial( destination_order.scalar_type() == source_order.scalar_type(), "edge_force_virial: destination_order and source_order must have the " "same int32, uint32, or int64 dtype"); + TORCH_CHECK( + edge_spin_gradient.numel() == 0 || + (edge_spin_gradient.is_cuda() && edge_spin_gradient.is_contiguous() && + edge_spin_gradient.device() == edge_gradient.device() && + edge_spin_gradient.sizes() == edge_gradient.sizes() && + edge_spin_gradient.scalar_type() == edge_gradient.scalar_type()), + "edge_force_virial: edge_spin_gradient must be empty or match " + "the gradient in device, layout, shape and dtype"); return assemble_force_virial(node_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, n_node_per_frame, - want_atom_virial); + edge_spin_gradient, want_atom_virial); } -std::tuple +std::tuple canonical_edge_force_virial(torch::Tensor edge_gradient, torch::Tensor edge_vec, torch::Tensor destination_row_ptr, torch::Tensor source_row_ptr, torch::Tensor source_order, torch::Tensor n_node_per_frame, + torch::Tensor edge_spin_gradient, c10::SymInt node_capacity, bool want_atom_virial) { const long node_count = node_capacity.expect_int(); @@ -483,13 +535,21 @@ canonical_edge_force_virial(torch::Tensor edge_gradient, source_row_ptr.numel() == node_count + 1, "canonical_edge_force_virial: row pointers must have N + 1 " "entries"); + TORCH_CHECK( + edge_spin_gradient.numel() == 0 || + (edge_spin_gradient.is_cuda() && edge_spin_gradient.is_contiguous() && + edge_spin_gradient.device() == edge_gradient.device() && + edge_spin_gradient.sizes() == edge_gradient.sizes() && + edge_spin_gradient.scalar_type() == edge_gradient.scalar_type()), + "canonical_edge_force_virial: edge_spin_gradient must be empty " + "or match the gradient in device, layout, shape and dtype"); auto edge_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); auto destination_order = torch::empty({0}, source_order.options()); return assemble_force_virial(node_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, n_node_per_frame, - want_atom_virial); + edge_spin_gradient, want_atom_virial); } // Per-frame total of a scalar carried on the node axis, the energy being the @@ -546,16 +606,19 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "edge_force_virial(Tensor edge_gradient, Tensor edge_vec, " "Tensor edge_index, Tensor edge_mask, Tensor destination_order, " "Tensor destination_row_ptr, Tensor source_order, Tensor source_row_ptr, " - "Tensor n_node_per_frame, SymInt node_capacity, " - "bool want_atom_virial) -> " - "(Tensor force, Tensor atom_virial, Tensor virial)"); + "Tensor n_node_per_frame, Tensor edge_spin_gradient, " + "SymInt node_capacity, bool want_atom_virial) -> " + "(Tensor force, Tensor atom_virial, Tensor virial, " + "Tensor magnetic_force)"); library.impl("edge_force_virial", torch::kCUDA, &edge_force_virial); library.def( "canonical_edge_force_virial(Tensor edge_gradient, Tensor edge_vec, " "Tensor destination_row_ptr, Tensor source_row_ptr, " - "Tensor source_order, Tensor n_node_per_frame, SymInt node_capacity, " + "Tensor source_order, Tensor n_node_per_frame, " + "Tensor edge_spin_gradient, SymInt node_capacity, " "bool want_atom_virial) -> " - "(Tensor force, Tensor atom_virial, Tensor virial)"); + "(Tensor force, Tensor atom_virial, Tensor virial, " + "Tensor magnetic_force)"); library.impl("canonical_edge_force_virial", torch::kCUDA, &canonical_edge_force_virial); library.def( diff --git a/source/op/pt/graph_ops.h b/source/op/pt/graph_ops.h index ee6de887dd..49fbf40fe4 100644 --- a/source/op/pt/graph_ops.h +++ b/source/op/pt/graph_ops.h @@ -166,17 +166,20 @@ void fitting_backward_range(cudaStream_t stream, float* d_x); // Scatter dE/d(edge_vec) into per-node force, per-frame virial and (optional) -// per-node virial. Returns (force (N, 3), atom_virial (N, 3, 3) or empty, -// virial (nf, 3, 3)). -std::tuple edge_force_virial( - torch::Tensor g_e, - torch::Tensor edge_vec, - torch::Tensor edge_index, - torch::Tensor edge_mask, - torch::Tensor destination_order, - torch::Tensor destination_row_ptr, - torch::Tensor source_order, - torch::Tensor source_row_ptr, - torch::Tensor n_node_per_frame, - c10::SymInt node_capacity, - bool want_atom_virial); +// per-node virial. A non-empty ``edge_spin_gradient`` adds the per-source total +// of the magnetic cotangent, which shares the source grouping the force +// reduction already walks. Returns (force (N, 3), atom_virial (N, 3, 3) or +// empty, virial (nf, 3, 3), magnetic_force (N, 3) or empty). +std::tuple +edge_force_virial(torch::Tensor g_e, + torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor destination_order, + torch::Tensor destination_row_ptr, + torch::Tensor source_order, + torch::Tensor source_row_ptr, + torch::Tensor n_node_per_frame, + torch::Tensor edge_spin_gradient, + c10::SymInt node_capacity, + bool want_atom_virial); diff --git a/source/tests/common/dpmodel/test_descriptor_dpa4c.py b/source/tests/common/dpmodel/test_descriptor_dpa4c.py index 8ba2220b98..46ab67a072 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa4c.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa4c.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import dataclasses +import math from typing import ( Any, ) @@ -105,8 +106,8 @@ def edge_features( return descriptor.build_edge_features( graph, atype_local, - *descriptor.pair_film.call(descriptor.type_embedding.call()), - ) + descriptor.pair_film.call(descriptor.type_embedding.call()), + )[:3] def moment_blocks( @@ -216,6 +217,8 @@ def test_moments_match_explicit_reference(self) -> None: amplitude, basis, envelope, + None, + None, dst, n_total, ) @@ -754,3 +757,607 @@ def test_automatic_profiles_and_output_dimensions( + 2 ) assert descriptor.get_dim_out() == expected_dim + + +# === Native spin === + +SPIN_COORD = np.array( + [ + [ + [0.0, 0.0, 0.0], + [1.4, 0.3, -0.2], + [-0.5, 1.2, 0.4], + [0.3, -0.7, 1.5], + [-0.9, -0.4, -1.0], + [1.1, -1.2, 0.6], + ] + ], + dtype=np.float64, +) +SPIN_ATYPE = np.array([[0, 1, 0, 1, 0, 1]], dtype=np.int64) + + +def make_spin_descriptor(**overrides: Any) -> DescrptDPA4C: + """Build a descriptor whose first atom type carries a magnetic moment.""" + config: dict[str, Any] = { + "rcut": 3.0, + "ntypes": 2, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 23, + "use_spin": [True, False], + } + config.update(overrides) + return DescrptDPA4C(**config) + + +def spin_reference_terms( + descriptor: DescrptDPA4C, + graph: Any, + atype: np.ndarray, + spin: np.ndarray, +) -> dict[str, np.ndarray]: + """Evaluate the two-body spin sums the readout is meant to reproduce. + + The reference reconstructs the edge stage from the descriptor's own + modules rather than from the moments, so a layout error cannot cancel + between the two sides of the comparison. + """ + masked = descriptor.spin.conditioned_spin(spin, atype) + source, destination = graph.edge_index[0], graph.edge_index[1] + edge_vec = graph.edge_vec + distance = np.sqrt( + np.sum(edge_vec * edge_vec, axis=-1, keepdims=True) + descriptor._EPS**2 + ) + direction = edge_vec / distance + center, neighbor = atype[destination], atype[source] + real = (center < descriptor.ntypes) & (neighbor < descriptor.ntypes) + envelope = descriptor.evaluate_cutoff_envelope(distance)[:, 0] * ( + graph.edge_mask & real + ) + radial = descriptor.radial_embedding.call(descriptor.radial_basis.call(distance)) + scale, shift, _, spin_scale, spin_shift = descriptor.pair_film.call( + descriptor.type_embedding.call() + ) + pair = center * (descriptor.ntypes + 1) + neighbor + channels = descriptor.spin_channels + weight = (radial[:, :channels] * spin_scale[pair] + spin_shift[pair]) * ( + envelope * envelope + )[:, None] + amplitude = (radial * scale[pair] + shift[pair]) * envelope[:, None] + + center_spin = masked[destination] + neighbor_spin = masked[source] + dot = np.sum(center_spin * neighbor_spin, axis=-1) + projection = np.sum(center_spin * direction, axis=-1) + neighbor_projection = np.sum(neighbor_spin * direction, axis=-1) + center_norm = np.sum(center_spin * center_spin, axis=-1) + neighbor_norm = np.sum(neighbor_spin * neighbor_spin, axis=-1) + + nodes = atype.shape[0] + + def reduce(values: np.ndarray) -> np.ndarray: + out = np.zeros((nodes, values.shape[1]), dtype=values.dtype) + np.add.at(out, destination, values) + return out + + return { + # Heisenberg exchange, one sum per spin channel. + "heisenberg": reduce(weight * dot[:, None]), + # Symmetric anisotropic two-ion exchange, one sum per spin channel. + "two_ion_anisotropy": reduce( + weight * (projection * neighbor_projection)[:, None] + ), + # Biquadratic exchange from the leading spin channel only, which is + # the one the neighbour quadrupole family reads. + "biquadratic": reduce( + weight[:, :1] + * (0.5 * (3.0 * dot * dot - center_norm * neighbor_norm))[:, None] + ), + # Single-ion anisotropy against every geometric degree-two channel. + "anisotropy": reduce( + amplitude[:, : descriptor.degree_channels[2]] + * envelope[:, None] + * (0.5 * (3.0 * projection * projection - center_norm))[:, None] + ), + } + + +def vector_gram_selectors(descriptor: DescrptDPA4C) -> dict[str, np.ndarray]: + """Return boolean column selectors of the joint degree-one spin Gram. + + The block holds the on-site moment, the ``spin_channels`` isotropic + neighbour channels and the ``spin_channels`` bond-projected ones, in that + order, so the entries of each physical interaction are addressed by their + channel coordinates rather than by a hard-coded offset. + """ + channels = descriptor.spin_channels + row, column = np.triu_indices(descriptor.spin.vector_width) + return { + "heisenberg": (row == 0) & (column >= 1) & (column <= channels), + "two_ion_anisotropy": (row == 0) & (column > channels), + "bond": (row > channels) | (column > channels), + } + + +class TestDPA4CSpin: + def setup_method(self) -> None: + self.descriptor = make_spin_descriptor() + rng = np.random.default_rng(5) + self.spin = rng.normal(size=(SPIN_ATYPE.size, 3)) + self.graph = neighbor_graph.build_neighbor_graph( + SPIN_COORD, + SPIN_ATYPE, + None, + self.descriptor.get_rcut(), + ) + self.atype = SPIN_ATYPE.reshape(-1) + + def evaluate(self, spin: np.ndarray | None) -> np.ndarray: + return self.descriptor.call_graph(self.graph, self.atype, spin=spin)[0] + + def spin_block(self, descriptor_output: np.ndarray) -> np.ndarray: + start = self.descriptor.readout.get_dim_out() + return descriptor_output[:, start : start + self.descriptor.spin.get_dim_out()] + + def test_axial_o3_invariance_including_reflections(self) -> None: + # Spin is an axial vector, so an improper transformation rotates it + # and flips its sign. Because every emitted coordinate has even spin + # order the descriptor is additionally invariant under the polar + # convention, which this test also pins. + rng = np.random.default_rng(11) + reference = self.evaluate(self.spin) + for determinant in (1.0, -1.0): + orthogonal, triangular = np.linalg.qr(rng.normal(size=(3, 3))) + orthogonal = orthogonal * np.sign(np.diag(triangular)) + if np.linalg.det(orthogonal) * determinant < 0.0: + orthogonal = -orthogonal + rotated = neighbor_graph.build_neighbor_graph( + SPIN_COORD @ orthogonal.T, + SPIN_ATYPE, + None, + self.descriptor.get_rcut(), + ) + for spin in ( + (self.spin @ orthogonal.T) * np.linalg.det(orthogonal), + self.spin @ orthogonal.T, + ): + output = self.descriptor.call_graph(rotated, self.atype, spin=spin)[0] + np.testing.assert_allclose(output, reference, atol=1e-12) + + def test_time_reversal_is_exact(self) -> None: + # Every spin family has even spin order, so a global moment flip is a + # bitwise symmetry rather than an approximate one. + np.testing.assert_array_equal( + self.evaluate(self.spin), + self.evaluate(-self.spin), + ) + + def test_non_magnetic_types_are_ignored_bitwise(self) -> None: + # The per-type gate is multiplicative, so a non-magnetic atom has no + # spin degree of freedom at any derivative order rather than merely a + # vanishing one. + polluted = self.spin.copy() + polluted[self.atype == 1] += 7.0 + np.testing.assert_array_equal( + self.evaluate(self.spin), + self.evaluate(polluted), + ) + + def test_spin_coordinates_vanish_with_the_moments(self) -> None: + block = self.spin_block(self.evaluate(np.zeros_like(self.spin))) + # Every family except the trailing magnetic coordination reads the + # spin value and is therefore exactly zero without moments. + channels = self.descriptor.spin_channels + np.testing.assert_array_equal( + block[:, :-channels], + np.zeros_like(block[:, :-channels]), + ) + assert np.any(block[:, -channels:] != 0.0) + + def test_a_missing_moment_is_rejected(self) -> None: + # A vanishing moment and an absent one are different states: the + # first is a demagnetized configuration, the second is a missing + # input whose magnetic force would be silently zero. + with pytest.raises(ValueError, match="requires a per-node magnetic"): + self.evaluate(None) + + def moment_divisor(self, spin: np.ndarray) -> np.ndarray: + """Return the neighborhood divisor the spin families are scaled by.""" + masked = self.descriptor.spin.conditioned_spin(spin, self.atype) + return self.descriptor.aggregate_moments( + *self.descriptor.build_edge_features( + self.graph, + self.atype, + self.descriptor.pair_film.call(self.descriptor.type_embedding.call()), + masked, + ), + self.descriptor.spin.onsite_payload(masked, self.atype), + self.graph.edge_index[1], + self.atype.shape[0], + )[1][:, 1] + + def test_two_body_terms_are_exactly_representable(self) -> None: + # The emitted invariants are not merely correlated with the physical + # sums: each is that sum times a known constant. + spin = self.descriptor.spin + reference = spin_reference_terms( + self.descriptor, + self.graph, + self.atype, + self.spin, + ) + block = self.spin_block(self.evaluate(self.spin)) + divisor = self.moment_divisor(self.spin) + vector_weight = spin.adam_spin_vector_weight[self.atype] + quadrupole_weight = spin.adam_spin_quadrupole_weight[self.atype] + + vector_gram = block[:, : spin.vector_gram_index.shape[0]] + selector = vector_gram_selectors(self.descriptor) + np.testing.assert_allclose( + vector_gram[:, selector["heisenberg"]], + math.sqrt(2.0) + * (vector_weight / divisor)[:, None] + * reference["heisenberg"], + atol=1e-12, + ) + + offset = spin.vector_gram_index.shape[0] + quadrupole_gram = block[ + :, offset : offset + spin.quadrupole_gram_index.shape[0] + ] + # The on-site self-term is not emitted, so the on-site x neighbour + # entry leads the quadrupole block. + np.testing.assert_allclose( + quadrupole_gram[:, 0:1], + math.sqrt(2.0) + * (quadrupole_weight / divisor)[:, None] + * reference["biquadratic"], + atol=1e-12, + ) + + offset += spin.quadrupole_gram_index.shape[0] + degree_two = self.descriptor.degree_channels[2] + cross = block[:, offset : offset + spin.quadrupole_width * degree_two].reshape( + -1, spin.quadrupole_width, degree_two + ) + np.testing.assert_allclose( + cross[:, 0, :], + (quadrupole_weight / divisor)[:, None] * reference["anisotropy"], + atol=1e-12, + ) + + def test_two_ion_anisotropy_is_exactly_representable(self) -> None: + # The pseudo-dipolar sum sum_j K(r_ij) (s_i.u_ij)(s_j.u_ij) is the + # interaction the bond-projected family exists for. It is emitted as + # an exact single sum, one entry per spin channel, so a linear + # readout spans an arbitrary K(r) inside the channel amplitudes. + spin = self.descriptor.spin + reference = spin_reference_terms( + self.descriptor, + self.graph, + self.atype, + self.spin, + ) + block = self.spin_block(self.evaluate(self.spin)) + divisor = self.moment_divisor(self.spin) + vector_weight = spin.adam_spin_vector_weight[self.atype] + + vector_gram = block[:, : spin.vector_gram_index.shape[0]] + selector = vector_gram_selectors(self.descriptor) + np.testing.assert_allclose( + vector_gram[:, selector["two_ion_anisotropy"]], + math.sqrt(2.0) + * (vector_weight / divisor)[:, None] + * reference["two_ion_anisotropy"], + atol=1e-12, + ) + # The two families are genuinely different observables and not one + # rescaled copy of the other. + assert not np.allclose( + reference["two_ion_anisotropy"], + reference["heisenberg"], + ) + + def test_output_width_matches_the_spin_layout(self) -> None: + spin = self.descriptor.spin + channels = self.descriptor.spin_channels + # The degree-one block carries the on-site moment plus the isotropic + # and bond-projected neighbour channels. The quadrupole Gram omits its + # on-site self-term, which the identity |B_2(s)|^2 = |s|^4 makes a + # function of the vector self-term. + width = 1 + 2 * channels + assert spin.vector_width == width + expected = ( + width * (width + 1) // 2 + + 2 + + 2 * self.descriptor.degree_channels[2] + + 2 * channels + ) + assert spin.get_dim_out() == expected + assert self.evaluate(self.spin).shape[1] == self.descriptor.get_dim_out() + assert ( + self.descriptor.get_dim_out() + == make_spin_descriptor(use_spin=None).get_dim_out() + expected + ) + + def test_serialization_roundtrip_preserves_spin(self) -> None: + self.descriptor.spin.set_spin_reference(np.array([1.7, 1.0, 1.0])) + restored = DescrptDPA4C.deserialize(self.descriptor.serialize()) + assert restored.use_spin == [True, False] + assert restored.supports_native_spin() + np.testing.assert_array_equal( + restored.spin.spin_reference, + self.descriptor.spin.spin_reference, + ) + np.testing.assert_allclose( + restored.call_graph(self.graph, self.atype, spin=self.spin)[0], + self.evaluate(self.spin), + atol=1e-14, + ) + + def test_sharing_rejects_a_different_spin_configuration(self) -> None: + with pytest.raises(ValueError, match="identical structural"): + self.descriptor.share_params(make_spin_descriptor(use_spin=None), 0) + + +#: Vertices of a regular tetrahedron. Every component has the same magnitude, +#: so the four squared norms agree bitwise and the four neighbours below share +#: one radial amplitude exactly rather than to rounding. +TETRAHEDRON = np.array( + [ + [1.0, 1.0, 1.0], + [1.0, -1.0, -1.0], + [-1.0, 1.0, -1.0], + [-1.0, -1.0, 1.0], + ], + dtype=np.float64, +) + + +def test_permuting_equidistant_moments_moves_the_descriptor() -> None: + """Relabelling equidistant neighbours must reach the output. + + Four identical neighbours sit at one distance from the centre, so every + radial spin family sees the same amplitude on every bond and is blind to + which moment sits on which bond. The physical pseudo-dipolar sum + ``sum_j (s_i.u_ij)(s_j.u_ij)`` is not blind to it, and the bond-projected + family is what carries that dependence into the readout. + """ + descriptor = make_spin_descriptor(ntypes=1, use_spin=[True]) + direction = TETRAHEDRON / np.linalg.norm(TETRAHEDRON, axis=-1, keepdims=True) + coord = np.concatenate([np.zeros((1, 3)), 1.5 * direction])[None] + atype = np.zeros((1, 5), dtype=np.int64) + graph = neighbor_graph.build_neighbor_graph( + coord, + atype, + None, + descriptor.get_rcut(), + ) + flat_atype = atype.reshape(-1) + + spin = np.random.default_rng(3).normal(size=(5, 3)) + relabelled = spin.copy() + relabelled[1:] = spin[[2, 3, 4, 1]] + + def pseudo_dipolar(moments: np.ndarray) -> float: + return float( + np.sum( + (moments[0] @ direction.T) + * np.einsum("jd,jd->j", moments[1:], direction) + ) + ) + + assert abs(pseudo_dipolar(spin) - pseudo_dipolar(relabelled)) > 1.0e-2 + + reference = descriptor.call_graph(graph, flat_atype, spin=spin)[0] + permuted = descriptor.call_graph(graph, flat_atype, spin=relabelled)[0] + + # The geometry is untouched, so the whole geometric block is unchanged. + geometry = descriptor.readout.get_dim_out() + np.testing.assert_allclose( + permuted[:, :geometry], + reference[:, :geometry], + atol=1e-13, + ) + + # Inside the spin block only the entries that touch the bond-projected + # channels may move: every other family reads the moments through a + # permutation-symmetric sum over the equidistant shell. + spin_width = descriptor.spin.get_dim_out() + before = reference[0, geometry : geometry + spin_width] + after = permuted[0, geometry : geometry + spin_width] + bond = np.zeros(spin_width, dtype=bool) + gram_width = descriptor.spin.vector_gram_index.shape[0] + bond[:gram_width] = vector_gram_selectors(descriptor)["bond"] + np.testing.assert_allclose(after[~bond], before[~bond], atol=1e-13) + assert np.abs(after[bond] - before[bond]).max() > 1.0e-2 + + +def test_periodic_cell_spin_matches_its_supercell() -> None: + """A periodic magnetic cell agrees with its own doubled cell. + + The graph builder folds every periodic image onto its local owner through + ``src = mapping[neighbor]``, so a cell narrower than the cutoff produces + edges whose source is the centre itself and edges that reach one owner + several times. Both are exercised here and by nothing else in this file. + """ + descriptor = make_spin_descriptor() + cell = 2.9 + coord = np.array([[[0.0, 0.0, 0.0], [1.45, 1.45, 0.2]]]) + atype = np.array([[0, 1]], dtype=np.int64) + box = np.array([[cell, 0.0, 0.0, 0.0, cell, 0.0, 0.0, 0.0, cell]]) + spin = np.random.default_rng(13).normal(size=(2, 3)) + + graph = neighbor_graph.build_neighbor_graph( + coord, + atype, + box, + descriptor.get_rcut(), + ) + source, destination = graph.edge_index[0], graph.edge_index[1] + assert np.any((source == destination) & graph.edge_mask) + + super_box = box.copy() + super_box[0, 0] = 2.0 * cell + super_graph = neighbor_graph.build_neighbor_graph( + np.concatenate([coord, coord + np.array([cell, 0.0, 0.0])], axis=1), + np.concatenate([atype, atype], axis=1), + super_box, + descriptor.get_rcut(), + ) + + reference = descriptor.call_graph(graph, atype.reshape(-1), spin=spin)[0] + doubled = descriptor.call_graph( + super_graph, + np.concatenate([atype, atype], axis=1).reshape(-1), + spin=np.concatenate([spin, spin], axis=0), + )[0] + np.testing.assert_allclose( + doubled, + np.concatenate([reference, reference], axis=0), + atol=1e-11, + ) + + +@pytest.mark.parametrize("channels", [8, 32, 128]) +def test_ordered_spin_tables_start_at_a_usable_scale(channels: int) -> None: + """Both spin tables must start near the scale of the geometric ones. + + The geometric heads are structurally anchored, so they start at a + root-mean-square of one and one quarter respectively. The spin heads may + not be anchored on a constant, because an exchange amplitude is signed, + and the descriptor calibration freezes a preconditioner at whatever scale + it measures. A spin table emerging from the bias-free trunk alone would + fix that preconditioner orders of magnitude below the block it belongs to. + """ + descriptor = make_spin_descriptor( + channels=channels, + ntypes=4, + n_radial=12, + use_spin=[True, True, False, False], + ) + scale, shift, _mixing, spin_scale, spin_shift = descriptor.pair_film.call( + descriptor.type_embedding.call() + ) + assert 0.9 <= float(np.sqrt(np.mean(np.square(scale)))) <= 1.1 + assert float(np.sqrt(np.mean(np.square(shift)))) > 0.1 + for table in (spin_scale, spin_shift): + assert 0.3 <= float(np.sqrt(np.mean(np.square(table)))) <= 0.7 + # The anchor fixes the magnitude of every channel, so the scale holds + # entry by entry rather than only in the aggregate. + assert float(np.abs(table).min()) >= 0.3 + assert float(np.abs(table).max()) <= 0.7 + + +def test_ordered_spin_tables_do_not_anchor_their_sign() -> None: + """The exchange amplitude of an ordered pair may take either sign. + + The geometric scale is anchored on the constant one and is therefore + strictly positive by construction. The spin tables carry no such bias: a + ferromagnetic and an antiferromagnetic pair are equally reachable from the + initialization, so both signs occur across seeds. + """ + signs: set[float] = set() + for seed in range(8): + descriptor = make_spin_descriptor(seed=seed) + _scale, _shift, _mixing, spin_scale, spin_shift = descriptor.pair_film.call( + descriptor.type_embedding.call() + ) + signs.update(np.sign(spin_scale).reshape(-1).tolist()) + signs.update(np.sign(spin_shift).reshape(-1).tolist()) + assert signs == {-1.0, 1.0} + + +#: Per-type moment scales of the calibration corpus. The second magnetic type +#: is rare and carries a much larger moment than the first. +CORPUS_MOMENT_SCALE = (1.0, 4.0, 0.0) + + +def magnetic_corpus(seed: int = 47) -> tuple[list[dict], np.ndarray]: + """Build a magnetic calibration corpus and its per-type moment scales. + + Returns + ------- + corpus + One sampled system carrying ``coord``, ``atype``, ``box`` and + ``spin``, in the packing ``compute_input_stats`` consumes. + reference + Independently computed per-type root-mean-square moment with shape + ``(ntypes + 1,)``. + """ + rng = np.random.default_rng(seed) + nframes, natoms, cell = 8, 12, 9.0 + # One rare magnetic atom of the second type per frame; the rest alternate + # between the abundant magnetic type and the non-magnetic one. + atype = np.tile(np.array([0, 0, 0, 2, 0, 2, 0, 0, 2, 0, 0, 1]), (nframes, 1)) + coord = rng.uniform(0.0, cell, size=(nframes, natoms, 3)) + box = np.tile(np.diag([cell, cell, cell]).reshape(1, 9), (nframes, 1)) + scale = np.take(np.asarray(CORPUS_MOMENT_SCALE), atype)[..., None] + spin = rng.normal(size=(nframes, natoms, 3)) * scale + + reference = np.ones(len(CORPUS_MOMENT_SCALE) + 1, dtype=np.float64) + magnitude = np.sum(np.square(spin), axis=-1) + for kind in range(len(CORPUS_MOMENT_SCALE)): + selected = magnitude[atype == kind] + if np.any(selected > 0.0): + reference[kind] = np.sqrt(np.mean(selected)) + corpus = [{"coord": coord, "atype": atype, "box": box, "spin": spin}] + return corpus, reference + + +def calibrated_descriptor(corpus: list[dict]) -> DescrptDPA4C: + """Return a three-type magnetic descriptor calibrated on ``corpus``.""" + descriptor = make_spin_descriptor(ntypes=3, use_spin=[True, True, False]) + descriptor.compute_input_stats(corpus) + return descriptor + + +def test_calibration_measures_the_per_type_reference_moment() -> None: + corpus, reference = magnetic_corpus() + descriptor = calibrated_descriptor(corpus) + np.testing.assert_allclose(descriptor.spin.spin_reference, reference, rtol=1e-12) + # The rare species keeps its own scale rather than the population one, so + # the conditioned moments of the two magnetic types land on one scale. + assert reference[1] / reference[0] > 3.0 + # A type observed only with a vanishing moment, and the padding row, keep + # the unit reference that leaves the raw spin untouched. + assert descriptor.spin.spin_reference[2] == 1.0 + assert descriptor.spin.spin_reference[3] == 1.0 + + +def test_calibration_conditions_every_spin_coordinate() -> None: + corpus, _reference = magnetic_corpus() + descriptor = calibrated_descriptor(corpus) + start = descriptor.readout.get_dim_out() + stddev = descriptor.stddev[start : start + descriptor.spin.get_dim_out()] + # Every spin coordinate activates on a magnetic corpus, so none of them + # falls back to the identity preconditioner, and none is driven to the + # extreme gain an unanchored spin table used to produce. + assert np.all(stddev > 0.0) + assert not np.any(stddev == 1.0) + assert float(stddev.max() / stddev.min()) < 1.0e4 + + +def test_calibration_accepts_either_moment_key() -> None: + corpus, _reference = magnetic_corpus() + renamed, _reference = magnetic_corpus() + renamed[0]["model_spin"] = renamed[0].pop("spin") + under_spin = calibrated_descriptor(corpus) + under_model_spin = calibrated_descriptor(renamed) + np.testing.assert_array_equal( + under_model_spin.spin.spin_reference, + under_spin.spin.spin_reference, + ) + np.testing.assert_array_equal(under_model_spin.stddev, under_spin.stddev) + + +def test_calibration_rejects_a_corpus_without_moments() -> None: + # A key mismatch would otherwise leave a unit reference magnitude and an + # identity preconditioner on most spin coordinates, with no error. + corpus, _reference = magnetic_corpus() + corpus[0].pop("spin") + with pytest.raises(ValueError, match="requires a per-node magnetic"): + calibrated_descriptor(corpus) diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 7036c5f811..d146aba9c3 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -51,6 +51,7 @@ p_examples / "spin" / "se_e2_a" / "input_torch.json", p_examples / "spin" / "dpa4" / "input.json", p_examples / "spin" / "dpa4" / "input-deepspin.json", + p_examples / "spin" / "dpa4c" / "input.json", p_examples / "dprc" / "normal" / "input.json", p_examples / "dprc" / "pairwise" / "input.json", p_examples / "dprc" / "generalized_force" / "input.json", diff --git a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py index 7280f66343..8adc39a421 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py @@ -1303,7 +1303,7 @@ def _assert_parity_vs_separate_ops(self, lmax: int) -> None: (g_e,) = torch.autograd.grad(e_atom.sum(), ev) # The fused operator assembles force / virial in the model compute # precision (fp32), so mirror that dtype in the reference scatter. - r_force, r_atom_vir, r_virial = edge_force_virial( + r_force, r_atom_vir, r_virial, _ = edge_force_virial( g_e.to(force.dtype), ev.detach().to(force.dtype), graph.edge_index, @@ -1313,6 +1313,7 @@ def _assert_parity_vs_separate_ops(self, lmax: int) -> None: graph.source_order, graph.source_row_ptr, graph.n_node, + ev.new_zeros(0, 3), n, True, ) @@ -1462,7 +1463,10 @@ def test_fused_energy_uses_owned_nodes_only(self) -> None: do_atomic_virial=True, ) assert fused is not None - energy, atom_energy, force, virial, atom_virial = fused + # A spin-free composition reports an empty magnetic force in the last + # position, which every implementation of the fused entry point emits. + energy, atom_energy, force, virial, atom_virial, magnetic = fused + assert magnetic.numel() == 0 with _CudaLevel("1"): edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) @@ -1477,7 +1481,7 @@ def test_fused_energy_uses_owned_nodes_only(self) -> None: :, None ] (edge_gradient,) = torch.autograd.grad(atom_energy_ref.sum(), edge_vec) - force_ref, atom_virial_ref, virial_ref = edge_force_virial( + force_ref, atom_virial_ref, virial_ref, _ = edge_force_virial( edge_gradient.to(force.dtype), edge_vec.detach().to(force.dtype), graph.edge_index, @@ -1487,6 +1491,7 @@ def test_fused_energy_uses_owned_nodes_only(self) -> None: graph.source_order, graph.source_row_ptr, graph.n_node, + edge_vec.new_zeros(0, 3), n_node, True, ) @@ -1622,7 +1627,7 @@ def _assert_parity_vs_separate_ops(self, lmax: int) -> None: grrg, _rot = dpa1_graph_compress(des, g2, atype, tebd) e_atom = fit.call_graph(grrg, atype)[fit.var_name] (g_e,) = torch.autograd.grad(e_atom.sum(), ev) - r_force, r_atom_vir, r_virial = edge_force_virial( + r_force, r_atom_vir, r_virial, _ = edge_force_virial( g_e.to(force.dtype), ev.detach().to(force.dtype), graph.edge_index, @@ -1632,6 +1637,7 @@ def _assert_parity_vs_separate_ops(self, lmax: int) -> None: graph.source_order, graph.source_row_ptr, graph.n_node, + ev.new_zeros(0, 3), n, True, ) @@ -1955,7 +1961,7 @@ def _fused( edge_force_virial, ) - return edge_force_virial( + force, atom_virial, virial, _ = edge_force_virial( g_e, edge_vec, edge_index, @@ -1965,9 +1971,11 @@ def _fused( src_order, src_row_ptr, n_node, + edge_vec.new_zeros(0, 3), total, True, ) + return force, atom_virial, virial def _assert_device_parity(self, device) -> None: args = self._random_graph(torch.device(device)) @@ -2047,6 +2055,7 @@ def test_compact_canonical_parity(self) -> None: source_order, source_row_ptr, n_node, + edge_vec.new_zeros(0, 3), total, True, ) @@ -2057,12 +2066,85 @@ def test_compact_canonical_parity(self) -> None: compact.source_row_ptr, compact.source_order, compact.n_node, + edge_vec.new_zeros(0, 3), total, True, ) for actual, expected in zip(canonical, generic, strict=True): torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + def test_magnetic_reduction_parity(self) -> None: + """The magnetic cotangent reduces onto its source nodes. + + The reduction rides the source loop of the force assembly, so it has to + agree with the reference on the grouping and on which edges a mask + removes. Both are checked against the CPU implementation on the same + graph. + """ + from deepmd.kernels.cuda.edge_force_virial import ( + edge_force_virial, + ) + + args = self._random_graph(torch.device("cuda")) + g_e, edge_vec, edge_index, mask = args[:4] + dst_order, dst_row_ptr, src_order, src_row_ptr, n_node, total = args[4:] + spin = torch.randn_like(edge_vec) + + def assemble(device: str) -> torch.Tensor: + move = lambda t: t.to(device) # noqa: E731 + return edge_force_virial( + move(g_e), + move(edge_vec), + move(edge_index), + move(mask), + move(dst_order), + move(dst_row_ptr), + move(src_order), + move(src_row_ptr), + move(n_node), + move(spin), + total, + True, + )[3] + + device_magnetic = assemble("cuda") + self.assertEqual(tuple(device_magnetic.shape), (total, 3)) + torch.testing.assert_close( + device_magnetic.cpu(), assemble("cpu"), atol=1e-10, rtol=1e-10 + ) + + # An absent cotangent yields an empty output and leaves the force alone. + force, _, _, empty = edge_force_virial( + g_e, + edge_vec, + edge_index, + mask, + dst_order, + dst_row_ptr, + src_order, + src_row_ptr, + n_node, + edge_vec.new_zeros(0, 3), + total, + True, + ) + self.assertEqual(empty.numel(), 0) + with_spin_force = edge_force_virial( + g_e, + edge_vec, + edge_index, + mask, + dst_order, + dst_row_ptr, + src_order, + src_row_ptr, + n_node, + spin, + total, + True, + )[0] + torch.testing.assert_close(force, with_spin_force) + def test_many_small_frames(self) -> None: """Frame reduction is valid beyond the CUDA grid-y limit.""" frame_count = 8192 diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py index a3488e1cea..e9f2bd3973 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -365,8 +365,8 @@ def test_amp_spans_the_edge_stage_and_stops_at_its_boundary( features = descriptor.build_edge_features( graph, atype_local, - *descriptor.pair_film.call(descriptor.type_embedding.call()), - ) + descriptor.pair_film.call(descriptor.type_embedding.call()), + )[:3] finally: for handle in handles: handle.remove() @@ -426,3 +426,149 @@ def test_coincident_edge_has_finite_third_derivative(self) -> None: atol=1e-12, rtol=0.0, ) + + +class TestDPA4CSpin: + """Torch-side contracts of the native spin branch.""" + + def setup_method(self) -> None: + self.descriptor = DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=8, + lmax=2, + n_radial=4, + precision="float64", + seed=23, + use_spin=[True, False], + ).to(env.DEVICE) + self.descriptor.eval() + generator = torch.Generator(device="cpu").manual_seed(5) + self.coord = ( + torch.randn(1, 6, 3, dtype=torch.float64, generator=generator).to( + env.DEVICE + ) + * 1.4 + ) + self.atype = torch.tensor( + [[0, 1, 0, 1, 0, 1]], dtype=torch.long, device=env.DEVICE + ) + self.spin = torch.randn(6, 3, dtype=torch.float64, generator=generator).to( + env.DEVICE + ) + from deepmd.dpmodel.utils.neighbor_graph import build_neighbor_graph + + self.graph = build_neighbor_graph(self.coord, self.atype, None, 3.0) + self.flat_atype = self.atype.reshape(-1) + + def probe(self, spin: torch.Tensor) -> torch.Tensor: + """Reduce the descriptor with fixed weights into a scalar probe.""" + output, _ = self.descriptor.call_graph(self.graph, self.flat_atype, spin=spin) + weights = torch.arange( + 1, output.shape[1] + 1, dtype=output.dtype, device=output.device + ) + return (output * weights).sum() + + def test_spin_arrays_are_optimizer_visible(self) -> None: + names = {name for name, _ in self.descriptor.named_parameters()} + assert "spin.adam_spin_vector_weight" in names + assert "spin.adam_spin_quadrupole_weight" in names + # The gate and the reference are state, not learned quantities. + buffers = {name for name, _ in self.descriptor.named_buffers()} + assert {"spin.spin_mask", "spin.spin_reference"} <= buffers + + def test_magnetic_force_matches_finite_differences(self) -> None: + # The magnetic force is the spin gradient of the energy, so the whole + # spin branch is validated against numerical differentiation rather + # than against stored values. + def descriptor_of_spin(spin: torch.Tensor) -> torch.Tensor: + return self.descriptor.call_graph(self.graph, self.flat_atype, spin=spin)[0] + + assert torch.autograd.gradcheck( + descriptor_of_spin, + (self.spin.clone().requires_grad_(True),), + ) + + def onsite_weights(self) -> list[torch.Tensor]: + """Return the two per-type on-site spin weights. + + These are the only parameters indexed by atom type that read the + moment value; the ordered spin tables are shared across types and are + also weighted by the moment-independent magnetic coordination family. + """ + return [ + self.descriptor.spin.adam_spin_vector_weight, + self.descriptor.spin.adam_spin_quadrupole_weight, + ] + + def test_non_magnetic_types_carry_no_magnetic_degree_of_freedom(self) -> None: + spin = self.spin.clone().requires_grad_(True) + (gradient,) = torch.autograd.grad(self.probe(spin), spin, create_graph=True) + assert torch.equal( + gradient[self.flat_atype == 1], + torch.zeros_like(gradient[self.flat_atype == 1]), + ) + # The force loss differentiates the magnetic force again, which probes + # the spin direction even where the value vanishes. A multiplicative + # gate is what keeps that second derivative exactly zero as well. + self.descriptor.zero_grad() + gradient.pow(2).sum().backward() + for weight in self.onsite_weights(): + assert weight.grad is not None + assert float(weight.grad[1].abs().max()) == 0.0 + + def test_zero_spin_leaves_no_magnetic_force(self) -> None: + spin = torch.zeros_like(self.spin).requires_grad_(True) + (gradient,) = torch.autograd.grad(self.probe(spin), spin) + assert torch.equal(gradient, torch.zeros_like(gradient)) + # Every route that reads the moment value is even in it, so the + # on-site weights stay dormant in the demagnetized limit. + self.descriptor.zero_grad() + self.probe(torch.zeros_like(self.spin)).backward() + for weight in self.onsite_weights(): + assert weight.grad is None or float(weight.grad.abs().max()) == 0.0 + + def test_a_missing_moment_is_rejected(self) -> None: + with pytest.raises(ValueError, match="requires a per-node magnetic"): + self.descriptor.call_graph(self.graph, self.flat_atype) + + def test_mixed_precision_never_engages(self) -> None: + # Spin families are quadratic in the moment and feed a fourth-order + # readout, so the branch stays in compute precision unconditionally. + for layer in self.descriptor.radial_embedding.layers: + assert layer.autocast_output is False + self.descriptor.use_amp = True + self.descriptor.use_amp_infer = True + self.descriptor._apply_autocast_policy() + for layer in self.descriptor.radial_embedding.layers: + assert layer.autocast_output is False + + def test_compression_covers_the_spin_families(self) -> None: + """The compiled operator carries spin, so eligibility ignores it. + + Every spin width follows the degree profile rather than a parameter of + its own, so a spin-conditioned descriptor is covered by the same + structural set as a spin-free one and its frozen tables are built + alongside the geometric caches. + """ + from deepmd.kernels.cuda.dpa4c.graph_compress import mega_eligible + + single = DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=8, + lmax=2, + n_radial=4, + precision="float32", + seed=23, + use_spin=[True, False], + ).to(env.DEVICE) + assert mega_eligible(single) + single.enable_compression(0.5) + assert single.compress + assert single.compress_spin_pair.shape == ( + (single.ntypes + 1) ** 2, + single.spin_channels, + 2, + ) + assert single.compress_spin_type.shape == (single.ntypes + 1, 4) diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py index ef5e5bf4b2..f710d8301c 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py @@ -3,6 +3,7 @@ import dataclasses +import numpy as np import pytest import torch @@ -19,6 +20,7 @@ descriptor_profile, dpa4c_graph_compress_energy_force, ensure_registered, + mega_eligible, op_available, ) from deepmd.pt.utils.nlist import ( @@ -116,12 +118,40 @@ def _arguments( artifacts["coupling_value"], artifacts["output_mean"], artifacts["output_inv_std"], + artifacts["spin_type"][:0], + artifacts["spin_pair"], + artifacts["spin_type"], bool(graph.destination_sorted), int(descriptor.lmax), *(float(value) for value in artifacts["info"]), ) +def _spin_free(arguments: tuple) -> tuple: + """Drop the native spin block for the spin-free CPU reference.""" + return (*arguments[:15], *arguments[18:]) + + +def _with_spin(arguments: tuple, spin: torch.Tensor) -> tuple: + """Place per-node magnetic moments in the operator's spin slot.""" + return (*arguments[:15], spin, *arguments[16:]) + + +def _assert_dispatched(actual: torch.Tensor, portable: torch.Tensor) -> None: + """Assert that a descriptor came from the compiled operator. + + The compressed path evaluates the radial embedding from a table, which + never reproduces the portable evaluation bit for bit. A run that fell back + to the portable code -- because the compression gate closed, or because + ``DP_CUDA_INFER`` left the operator disabled -- reproduces it exactly, and + would otherwise be compared with itself and pass every tolerance. + """ + assert not torch.equal(actual, portable), ( + "the compressed descriptor is bitwise identical to the portable one, " + "so the compiled operator did not run" + ) + + @_GPU @pytest.mark.parametrize("channels", [8, 16, 32, 64, 128]) @pytest.mark.parametrize("canonical", [False, True]) @@ -150,7 +180,7 @@ def test_forward_backward_parity(channels: int, canonical: bool) -> None: (gradient,) = torch.autograd.grad((output * cotangent).sum(), edge_vec) reference_edge = graph.edge_vec.detach().clone().requires_grad_(True) - reference = _cpu_descriptor(reference_edge, *arguments) + reference = _cpu_descriptor(reference_edge, *_spin_free(arguments)) (reference_gradient,) = torch.autograd.grad( (reference * cotangent).sum(), reference_edge, @@ -170,11 +200,15 @@ def test_backward_tail_node_groups(channels: int) -> None: descriptor = _build_descriptor(channels) graph, atype = _build_graph(descriptor, canonical=True, node_count=23) arguments = _arguments(descriptor, graph, atype) + # The cotangent fixes which element sits closest to the tolerance, so + # drawing it from the unseeded global stream would make the outcome vary + # between processes. cotangent = torch.randn( atype.shape[0], descriptor.get_dim_out(), dtype=torch.float32, device="cuda", + generator=torch.Generator(device="cuda").manual_seed(37), ) edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) @@ -185,7 +219,7 @@ def test_backward_tail_node_groups(channels: int) -> None: (gradient,) = torch.autograd.grad((output * cotangent).sum(), edge_vec) reference_edge = graph.edge_vec.detach().clone().requires_grad_(True) - reference = _cpu_descriptor(reference_edge, *arguments) + reference = _cpu_descriptor(reference_edge, *_spin_free(arguments)) (reference_gradient,) = torch.autograd.grad( (reference * cotangent).sum(), reference_edge, @@ -435,7 +469,7 @@ def test_supported_surface_parity( ).reshape_as(output) (gradient,) = torch.autograd.grad((output * cotangent).sum(), edge_vec) reference_edge = graph.edge_vec.detach().clone().requires_grad_(True) - reference_value = _cpu_descriptor(reference_edge, *arguments) + reference_value = _cpu_descriptor(reference_edge, *_spin_free(arguments)) (reference_gradient,) = torch.autograd.grad( (reference_value * cotangent).sum(), reference_edge, @@ -780,8 +814,8 @@ def test_compact_canonical_parity( graph.edge_index[0].to(index_dtype), graph.destination_row_ptr, atype, - *arguments[5:15], - *arguments[16:], + *arguments[5:18], + *arguments[19:], ) compact_output, compact_state = torch.ops.deepmd.dpa4c_canonical_compress( graph.edge_vec, @@ -825,8 +859,8 @@ def test_compact_inplace_backward_reuses_state(channels: int) -> None: graph.edge_index[0].to(torch.uint32), graph.destination_row_ptr, atype, - *arguments[5:15], - *arguments[16:], + *arguments[5:18], + *arguments[19:], ) ensure_canonical_registered() output, state = torch.ops.deepmd.dpa4c_canonical_compress( @@ -909,7 +943,7 @@ def test_fused_energy_force_parity( ) atom_energy = fitting.call_graph(node_descriptor, atype)[fitting.var_name] (edge_gradient,) = torch.autograd.grad(atom_energy.sum(), edge_vec) - force, atom_virial, virial = edge_force_virial( + force, atom_virial, virial, _ = edge_force_virial( edge_gradient, edge_vec.detach(), graph.edge_index, @@ -919,6 +953,7 @@ def test_fused_energy_force_parity( graph.source_order, graph.source_row_ptr, graph.n_node, + edge_vec.new_zeros(0, 3), atype.shape[0], True, ) @@ -1057,3 +1092,265 @@ def run() -> tuple[torch.Tensor, ...]: monkeypatch.setenv("DP_NODE_TILE", tile) for actual, expected in zip(run(), reference, strict=True): torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-6) + + +def _build_spin_descriptor( + channels: int, + lmax: int = 2, + radial_modes: int = 0, +) -> DescrptDPA4C: + """Return a spin-conditioned descriptor with a non-unit reference moment. + + A reference magnitude other than one makes the conditioning factor visible + in the magnetic gradient, so a missing chain factor shows up as a constant + ratio rather than cancelling. + """ + descriptor = ( + DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=channels, + lmax=lmax, + n_radial=8, + radial_modes=radial_modes, + precision="float32", + seed=17, + use_spin=[True, False], + ) + .cuda() + .eval() + ) + descriptor.spin.set_spin_reference(np.array([1.7, 1.0, 1.0])) + return descriptor + + +@_GPU +@pytest.mark.parametrize( + ("channels", "lmax", "radial_modes"), + [(8, 2, 0), (16, 2, 0), (32, 2, 4), (64, 3, 0), (128, 3, 8), (32, 4, 0)], +) +def test_spin_compressed_matches_portable( + monkeypatch: pytest.MonkeyPatch, + channels: int, + lmax: int, + radial_modes: int, +) -> None: + """Descriptor, coordinate gradient and magnetic force all match. + + The portable path is the oracle for the spin families. Probing every + output column at once keeps the geometric and the spin blocks in the same + comparison, because the two are coupled through the shared normalizer and + through the cross Gram. The cotangent stays bounded: weighting columns by + their index inflates the gradient magnitude at the widest profile until + fp32 tabulation noise alone exceeds the tolerance. + """ + descriptor = _build_spin_descriptor(channels, lmax, radial_modes) + assert mega_eligible(descriptor) + graph, atype = _build_graph(descriptor, canonical=True) + generator = torch.Generator(device="cuda").manual_seed(11) + spin = torch.randn( + atype.shape[0], + 3, + dtype=torch.float32, + device="cuda", + generator=generator, + ) + + def run() -> tuple[torch.Tensor, ...]: + moment = spin.detach().clone().requires_grad_(True) + edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) + output, _ = descriptor.call_graph( + dataclasses.replace(graph, edge_vec=edge_vec), + atype, + spin=moment, + ) + cotangent = torch.linspace( + -0.7, + 1.3, + output.numel(), + dtype=output.dtype, + device=output.device, + ).reshape_as(output) + gradients = torch.autograd.grad( + (output * cotangent).sum(), + [edge_vec, moment], + ) + return (output, *gradients) + + reference = run() + descriptor.enable_compression(0.5) + monkeypatch.setenv("DP_CUDA_INFER", "1") + assert descriptor.compress + actual = run() + _assert_dispatched(actual[0], reference[0]) + # Tabulation error reaches the gradients through the table derivative, so + # they carry the wider tolerance the geometric backward already uses. + tolerances = ((3e-5, 3e-5), (8e-6, 1e-4), (8e-6, 1e-4)) + for (atol, rtol), value, expected in zip( + tolerances, actual, reference, strict=True + ): + torch.testing.assert_close(value, expected, atol=atol, rtol=rtol) + + +@_GPU +@pytest.mark.parametrize("family", ["neighbour", "onsite"]) +def test_spin_magnetic_force_splits_into_onsite_and_neighbour( + monkeypatch: pytest.MonkeyPatch, + family: str, +) -> None: + """Each half of the magnetic force is correct on its own. + + The on-site half closes inside the node kernel while the neighbour half is + emitted per edge and reduced onto source nodes, so they fail + independently. Silencing one at a time keeps a fault in either from being + masked by the other's magnitude. + """ + descriptor = _build_spin_descriptor(16) + with torch.no_grad(): + if family == "neighbour": + descriptor.spin.adam_spin_vector_weight.zero_() + descriptor.spin.adam_spin_quadrupole_weight.zero_() + else: + geometric = descriptor.channels * (2 + descriptor.radial_modes) + for parameter in descriptor.pair_film.parameters(): + if parameter.dim() == 2 and parameter.shape[1] > geometric: + parameter[:, geometric:] = 0.0 + elif parameter.dim() == 1 and parameter.shape[0] > geometric: + parameter[geometric:] = 0.0 + graph, atype = _build_graph(descriptor, canonical=True) + generator = torch.Generator(device="cuda").manual_seed(11) + spin = torch.randn( + atype.shape[0], + 3, + dtype=torch.float32, + device="cuda", + generator=generator, + ) + + def run() -> tuple[torch.Tensor, torch.Tensor]: + moment = spin.detach().clone().requires_grad_(True) + output, _ = descriptor.call_graph(graph, atype, spin=moment) + cotangent = torch.linspace( + -0.7, + 1.3, + output.numel(), + dtype=output.dtype, + device=output.device, + ).reshape_as(output) + gradient = torch.autograd.grad((output * cotangent).sum(), moment)[0] + return output, gradient + + reference_output, reference_gradient = run() + descriptor.enable_compression(0.5) + monkeypatch.setenv("DP_CUDA_INFER", "1") + output, gradient = run() + _assert_dispatched(output, reference_output) + assert torch.count_nonzero(gradient).item() > 0 + torch.testing.assert_close(gradient, reference_gradient, atol=8e-6, rtol=1e-4) + + +@_GPU +def test_spin_bond_family_couples_the_kernel_to_the_edge_direction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The compiled spin path reads the edge direction, and only through it. + + Rotating the neighbourhood while holding the moments fixed leaves every + spin family except the bond-projected one invariant, so a kernel that + omitted that family would return an unchanged readout. Rotating the + moments alongside the geometry restores the invariance, which is what + separates a genuine bond coupling from a defect in the geometric block. + """ + descriptor = _build_spin_descriptor(16) + graph, atype = _build_graph(descriptor, canonical=True) + generator = torch.Generator(device="cuda").manual_seed(11) + spin = torch.randn( + atype.shape[0], + 3, + dtype=torch.float32, + device="cuda", + generator=generator, + ) + angle = 0.4 + rotation = torch.tensor( + [ + [float(np.cos(angle)), -float(np.sin(angle)), 0.0], + [float(np.sin(angle)), float(np.cos(angle)), 0.0], + [0.0, 0.0, 1.0], + ], + dtype=torch.float32, + device="cuda", + ) + + def run(edge_vec: torch.Tensor, moment: torch.Tensor) -> torch.Tensor: + output, _ = descriptor.call_graph( + dataclasses.replace(graph, edge_vec=edge_vec), + atype, + spin=moment, + ) + return output + + portable = run(graph.edge_vec, spin) + descriptor.enable_compression(0.5) + monkeypatch.setenv("DP_CUDA_INFER", "1") + upright = run(graph.edge_vec, spin) + _assert_dispatched(upright, portable) + rotated = run(graph.edge_vec @ rotation.T, spin) + assert not torch.allclose(rotated, upright, atol=1e-4) + covariant = run(graph.edge_vec @ rotation.T, spin @ rotation.T) + torch.testing.assert_close(covariant, upright, atol=3e-5, rtol=3e-5) + + +@_GPU +@pytest.mark.parametrize("spin_conditioned", [False, True]) +def test_backward_operator_satisfies_its_schema(spin_conditioned: bool) -> None: + """The backward operator declares three independent results. + + All three are unannotated, so any two of them sharing storage would be an + alias the schema does not describe, which is undefined under + functionalization. ``opcheck`` decides that mechanically, including on the + spin-free path where two of the three are absent and an allocation shared + between them would otherwise go unnoticed. + """ + descriptor = _build_spin_descriptor(8) if spin_conditioned else _build_descriptor(8) + graph, atype = _build_graph(descriptor, canonical=True) + arguments = _arguments(descriptor, graph, atype) + if spin_conditioned: + arguments = _with_spin( + arguments, + torch.randn(atype.shape[0], 3, dtype=torch.float32, device="cuda"), + ) + output, state = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec, + *arguments, + ) + torch.library.opcheck( + torch.ops.deepmd.dpa4c_graph_compress_backward.default, + (torch.ones_like(output), state, graph.edge_vec, *arguments), + ) + + +@_GPU +def test_registered_autograd_refuses_the_magnetic_moment() -> None: + """A direct operator call cannot silently lose the magnetic force. + + The operator emits that cotangent in two pieces and the per-edge piece is + reduced onto source nodes through the source CSR, which the schema does + not carry. The registration therefore cannot close the magnetic force, and + refusing is the only alternative to reporting a vanishing one. + """ + descriptor = _build_spin_descriptor(8) + graph, atype = _build_graph(descriptor, canonical=True) + spin = torch.randn( + atype.shape[0], + 3, + dtype=torch.float32, + device="cuda", + requires_grad=True, + ) + output, _state = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec, + *_with_spin(_arguments(descriptor, graph, atype), spin), + ) + with pytest.raises(RuntimeError, match="cannot differentiate its magnetic"): + output.sum().backward() diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py index 35f82b3a95..8e0b1447de 100644 --- a/source/tests/pt_expt/model/test_dpa4_native_spin.py +++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py @@ -1628,12 +1628,12 @@ def test_metadata_carries_pair_exclude_types(self) -> None: ) meta = _collect_metadata( - self._generic_model([[0, 1]]), is_spin=True, lower_kind="graph" + self._generic_model([[0, 1]]), spin_scheme="native", lower_kind="graph" ) assert meta["pair_exclude_types"] == [[0, 1]] base_meta = _collect_metadata( - self._generic_model(None), is_spin=True, lower_kind="graph" + self._generic_model(None), spin_scheme="native", lower_kind="graph" ) assert base_meta["pair_exclude_types"] == [] diff --git a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py index 002f870382..376b679c31 100644 --- a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py @@ -52,16 +52,23 @@ def _compressed_config(channels: int = 8) -> dict: fitting["neuron"] = [32, 32] fitting["activation_function"] = "silu" fitting["precision"] = "float32" + # The fused fitting operator has no per-layer timestep, and the shipped + # DPA4C grades do not use one either, so the compact canonical path is + # only reachable without it. + fitting["resnet_dt"] = False return config -def _run_graph(model: torch.nn.Module) -> dict[str, torch.Tensor]: +def _run_graph( + model: torch.nn.Module, + dtype: torch.dtype = torch.float64, +) -> dict[str, torch.Tensor]: sample = build_synthetic_graph_inputs( model, e_max=None, nframes=2, nloc=7, - dtype=torch.float64, + dtype=dtype, device=env.DEVICE, ) ( @@ -198,10 +205,15 @@ def test_compressed_level_two_matches_autograd( ) -> None: model = get_model(_compressed_config(channels)).to(env.DEVICE).eval() model.get_descriptor().enable_compression(min_nbor_dist=0.5) + # A compressed model is deployed on single-precision edge vectors. The + # level-two composition assembles force and virial in that precision + # throughout, whereas the autograd lower returns them in the precision of + # its edge leaf, so a double-precision sample would compare two different + # element types rather than two code paths. monkeypatch.setenv("DP_CUDA_INFER", "1") - reference = _run_graph(model) + reference = _run_graph(model, dtype=torch.float32) monkeypatch.setenv("DP_CUDA_INFER", "2") - actual = _run_graph(model) + actual = _run_graph(model, dtype=torch.float32) for key in ( "energy", "energy_redu", @@ -210,3 +222,104 @@ def test_compressed_level_two_matches_autograd( "energy_derv_c_redu", ): torch.testing.assert_close(actual[key], reference[key]) + + +def _spin_config(channels: int = 16) -> dict: + config = _compressed_config(channels) + config["descriptor"]["use_spin"] = [True, False] + return config + + +def _spin_sample(model: torch.nn.Module) -> tuple: + """Return a canonical graph, its flat types and a per-node moment.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + attach_edge_csr, + build_neighbor_graph, + ) + + nodes = 24 + generator = torch.Generator(device=env.DEVICE).manual_seed(31) + coord = 5.0 * torch.rand( + 1, nodes, 3, dtype=torch.float32, device=env.DEVICE, generator=generator + ) + atype = torch.arange(nodes, device=env.DEVICE).reshape(1, -1) % 2 + graph = attach_edge_csr( + build_neighbor_graph(coord, atype, None, model.get_rcut()), nodes + ) + spin = torch.randn( + nodes, 3, dtype=torch.float32, device=env.DEVICE, generator=generator + ) + return graph, atype.reshape(-1), spin + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="the fused spin path is CUDA only" +) +def test_compressed_spin_lowers_match_autograd( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Both fused lowers reproduce the autograd magnetic force. + + The magnetic force is assembled from two halves that no other output + exercises: the on-site gradient closes inside the node kernel, and the + neighbour half is emitted per edge and reduced onto source nodes. Checking + it against the autograd lower covers the fused generic composition and the + compact canonical deployment path in the same comparison. + """ + import numpy as np + + model = get_model(_spin_config()).to(env.DEVICE).eval() + descriptor = model.get_descriptor() + descriptor.spin.set_spin_reference(np.array([1.7, 1.0, 1.0])) + graph, atype, spin = _spin_sample(model) + descriptor.enable_compression(min_nbor_dist=0.5) + + def lower() -> dict[str, torch.Tensor]: + return model.forward_common_lower_graph( + atype, + graph.n_node, + graph.n_node.clone(), + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + destination_order=graph.destination_order, + destination_row_ptr=graph.destination_row_ptr, + source_order=graph.source_order, + source_row_ptr=graph.source_row_ptr, + destination_sorted=bool(graph.destination_sorted), + spin=spin, + ) + + monkeypatch.setenv("DP_CUDA_INFER", "1") + reference = lower() + assert "energy_derv_r_mag" in reference + monkeypatch.setenv("DP_CUDA_INFER", "2") + fused = lower() + for key in ("energy", "energy_redu", "energy_derv_r", "energy_derv_r_mag"): + torch.testing.assert_close(fused[key], reference[key], atol=8e-6, rtol=1e-4) + + physical = int(graph.destination_row_ptr[-1]) + canonical = model.forward_lower_canonical_graph( + atype, + graph.n_node, + graph.n_node.clone(), + graph.edge_index[0][:physical].to(torch.uint32).contiguous(), + graph.edge_vec[:physical].contiguous(), + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order[:physical].to(torch.uint32).contiguous(), + do_atomic_virial=False, + spin=spin, + ) + torch.testing.assert_close( + canonical["force_mag"].reshape(-1, 3), + reference["energy_derv_r_mag"].reshape(-1, 3), + atol=8e-6, + rtol=1e-4, + ) + torch.testing.assert_close( + canonical["force"].reshape(-1, 3), + reference["energy_derv_r"].reshape(-1, 3), + atol=8e-6, + rtol=1e-4, + ) diff --git a/source/tests/pt_expt/model/test_zbl_bridging.py b/source/tests/pt_expt/model/test_zbl_bridging.py index 3fa45a1c9c..28ec95b819 100644 --- a/source/tests/pt_expt/model/test_zbl_bridging.py +++ b/source/tests/pt_expt/model/test_zbl_bridging.py @@ -702,7 +702,7 @@ def test_bridged_metadata_carries_charge_spin_dim(tmp_path) -> None: config = copy.deepcopy(ZBL_CONFIG) config["descriptor"]["add_chg_spin_ebd"] = True model = get_model(config).to(torch.device("cpu")).eval() - meta = _collect_metadata(model, is_spin=False, lower_kind="graph") + meta = _collect_metadata(model, lower_kind="graph") assert meta["dim_chg_spin"] > 0, ( "the bridged model's metadata dropped charge_spin; the exported " "artifact would silently ignore the FiLM conditioning" @@ -713,7 +713,7 @@ def test_bridged_metadata_carries_charge_spin_dim(tmp_path) -> None: for key in ("bridging_method", "bridging_r_inner", "bridging_r_outer"): plain.pop(key, None) plain_model = get_model(plain).to(torch.device("cpu")).eval() - plain_meta = _collect_metadata(plain_model, is_spin=False, lower_kind="graph") + plain_meta = _collect_metadata(plain_model, lower_kind="graph") assert meta["dim_chg_spin"] == plain_meta["dim_chg_spin"] From 5f697a4d6d6a96b45e0e31b13177f1dfcf66afd4 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 30 Jul 2026 10:46:40 +0800 Subject: [PATCH 03/10] feat(dpa4c): add runtime charge-state conditioning and safe graph folding Condition DPA4C on frame charge and multiplicity while keeping it independent of per-atom native spin. - inject charge-state features into the type and ordered-pair routes - rebuild compressed constants once per runtime state and expose the setting across evaluation and LAMMPS interfaces - preserve unconditioned behavior and validate portable/compressed parity - reject graph folding without valid ghost-owner mappings instead of silently dropping halo edges --- deepmd/dpmodel/descriptor/dpa4_nn/mlp.py | 43 ++ deepmd/dpmodel/descriptor/dpa4c.py | 332 ++++++++++++++-- .../dpmodel/descriptor/dpa4c_nn/__init__.py | 8 + .../descriptor/dpa4c_nn/charge_state.py | 374 ++++++++++++++++++ .../dpmodel/descriptor/dpa4c_nn/pair_film.py | 134 +++++-- deepmd/dpmodel/descriptor/dpa4c_nn/spin.py | 15 +- deepmd/kernels/cuda/dpa4c/graph_compress.py | 293 ++++++++++---- deepmd/pt_expt/descriptor/dpa4c.py | 67 +++- deepmd/pt_expt/infer/deep_eval.py | 325 ++++++++++++++- deepmd/pt_expt/utils/serialization.py | 198 ++++++++++ deepmd/utils/argcheck.py | 28 ++ source/api_c/include/c_api.h | 58 ++- source/api_c/include/deepmd.hpp | 52 +++ source/api_c/src/c_api.cc | 48 +++ source/api_cc/include/DeepPot.h | 39 ++ source/api_cc/include/DeepPotPTExpt.h | 56 ++- source/api_cc/include/DeepSpin.h | 39 ++ source/api_cc/include/DeepSpinPTExpt.h | 11 + source/api_cc/include/NativeSpinPTExpt.h | 153 ++++++- source/api_cc/include/commonPT.h | 25 +- source/api_cc/src/DeepPot.cc | 4 + source/api_cc/src/DeepPotPTExpt.cc | 222 ++++++----- source/api_cc/src/DeepSpin.cc | 4 + source/api_cc/src/DeepSpinPTExpt.cc | 109 ++--- source/api_cc/src/NativeSpinPTExpt.cc | 192 ++++++++- source/api_cc/src/commonPTExpt.h | 215 +++++++++- source/lmp/pair_deepmd.cpp | 16 + source/lmp/pair_deepmd_kokkos.cpp | 10 +- source/lmp/pair_deepspin.cpp | 16 + source/lmp/pair_dpa4spin.cpp | 66 +++- source/lmp/pair_dpa4spin.h | 9 +- .../common/dpmodel/test_descriptor_dpa4c.py | 347 +++++++++++++++- source/tests/pt_expt/descriptor/test_dpa4c.py | 10 +- .../pt_expt/descriptor/test_dpa4c_cuda.py | 113 ++++++ .../pt_expt/model/test_dpa4c_graph_lower.py | 111 ++++++ 35 files changed, 3406 insertions(+), 336 deletions(-) create mode 100644 deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/mlp.py b/deepmd/dpmodel/descriptor/dpa4_nn/mlp.py index da90d08117..62e59cb88b 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/mlp.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/mlp.py @@ -155,6 +155,49 @@ def call_hidden(self, inputs: Any) -> Any: output = self.activation(layer(output)) return output + def call_hidden_affine(self, inputs: Any) -> Any: + """Apply the first hidden affine map without its activation. + + The affine map is linear and bias free, so an additive shift of the + input appears here as an additive shift of the pre-activation. A + caller that evaluates the same trunk under several such shifts can + therefore share this projection and add each shift afterwards, + instead of duplicating the input over the shift axis. + + Parameters + ---------- + inputs + Input with shape ``(..., mlp_layers[0])``. + + Returns + ------- + Any + Pre-activation with shape ``(..., 2 * mlp_layers[1])``. + """ + return self.layers[0](inputs) + + def call_from_hidden_affine(self, pre_activation: Any) -> Any: + """Complete the MLP from the first hidden affine pre-activation. + + Composing this with :meth:`call_hidden_affine` reproduces + :meth:`call` exactly, for any number of hidden layers. + + Parameters + ---------- + pre_activation + Pre-activation with shape ``(..., 2 * mlp_layers[1])``, as + returned by :meth:`call_hidden_affine`. + + Returns + ------- + Any + Output with shape ``(..., mlp_layers[-1])``. + """ + output = self.activation(pre_activation) + for layer in self.layers[1:-1]: + output = self.activation(layer(output)) + return self.call_output(output) + def call_output(self, hidden: Any) -> Any: """Apply the final scaled linear projection to a latent state. diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index f854259c9f..f413adf126 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -73,15 +73,18 @@ resolve_swiglu_hidden_width, ) from .dpa4c_nn import ( + ChargeStateEmbedding, InvariantReadout, OrderedPairFiLM, SpinChannels, build_angular_basis, build_moment_indices, + canonicalize_charge_spin, degree_offsets, derive_bispectrum_ranks, derive_degree_channels, derive_spin_channels, + validate_charge_state, ) if TYPE_CHECKING: @@ -174,6 +177,14 @@ class DescrptDPA4C(NativeOP, BaseDescriptor): given, the descriptor conditions on a per-node spin vector and declares :meth:`supports_native_spin`. ``None`` reproduces the spin-free descriptor exactly. + add_chg_spin_ebd + Whether to condition on the frame-level total charge and spin + multiplicity. This is unrelated to ``use_spin``, which carries a + per-atom magnetic moment. + default_chg_spin + Fallback ``[charge, multiplicity]`` used when a caller supplies no + explicit condition. Compression bakes this value, so a deployed + artifact requires it. spin Reserved for descriptor API compatibility; only ``None`` is supported. Native spin is configured through ``use_spin``. @@ -233,6 +244,8 @@ def __init__( type_map: list[str] | None = None, seed: int | list[int] | None = None, use_spin: list[bool] | None = None, + add_chg_spin_ebd: bool = False, + default_chg_spin: list[float] | None = None, spin: None = None, ) -> None: # === Step 1. Validate the public architecture contract === @@ -253,6 +266,11 @@ def __init__( or radial_modes < 0 ): raise ValueError("`radial_modes` must be a non-negative integer.") + default_chg_spin = ( + None + if default_chg_spin is None + else validate_charge_state(default_chg_spin) + ) # `channels` and `lmax` are validated inside the profile derivation, # which owns their supported sets. degree_channels = derive_degree_channels(channels, lmax) @@ -274,6 +292,8 @@ def __init__( self.type_map = type_map self.seed = seed self.use_spin = None if use_spin is None else [bool(flag) for flag in use_spin] + self.add_chg_spin_ebd = bool(add_chg_spin_ebd) + self.default_chg_spin = default_chg_spin # The spin branch reads the leading channels of the shared radial map, # so its width is derived rather than exposed. self.spin_channels = ( @@ -355,6 +375,17 @@ def __init__( seed=child_seed(seed, 5), ) ) + self.charge_spin_embedding = ( + ChargeStateEmbedding( + self.channels, + self.pair_film.pair_hidden_width, + precision=self.precision, + trainable=self.trainable, + seed=child_seed(seed, 6), + ) + if self.add_chg_spin_ebd + else None + ) # === Step 4. Lay out the flat moment payload === # Degree zero owns the leading `channels` entries of the flat layout, @@ -385,6 +416,7 @@ def call_graph( type_embedding: Array | None = None, comm_dict: dict | None = None, spin: Array | None = None, + charge_spin: Array | None = None, ) -> tuple[Array, None]: """Evaluate DPA4C on a flat neighbor graph. @@ -410,6 +442,11 @@ def call_graph( axis as ``atype``, including ghost and padding rows. Mandatory when the descriptor is configured with ``use_spin`` and ignored otherwise. + charge_spin + Frame-level total charge and spin multiplicity with shape + ``(nf, 2)``, or a single pair broadcast over the frames. Read + only when the descriptor is configured with ``add_chg_spin_ebd``, + which then falls back to ``default_chg_spin`` if this is absent. Returns ------- @@ -422,7 +459,8 @@ def call_graph( Raises ------ ValueError - If the descriptor is spin conditioned and ``spin`` is absent. + If the descriptor is spin conditioned and ``spin`` is absent, or + charge conditioned with neither ``charge_spin`` nor a default. """ del comm_dict # === Step 1. Resolve type features and compute precision === @@ -438,7 +476,15 @@ def call_graph( ) # === Step 2. Evaluate the graph-native equations === - descriptor, _ = self.evaluate_graph(graph, atype, type_embedding, spin) + descriptor, _ = self.evaluate_graph( + graph, + atype, + type_embedding, + spin, + self.require_charge_spin( + charge_spin, graph.n_node.shape[0], graph.edge_vec + ), + ) # === Step 3. Restore the graph input dtype === if descriptor.dtype != in_dtype: @@ -480,6 +526,38 @@ def require_spin(self, spin: Array | None) -> Array: ) return spin + def require_charge_spin( + self, + charge_spin: Array | None, + nf: int, + ref: Array, + ) -> Array | None: + """Resolve the frame condition a charge-conditioned descriptor needs. + + Parameters + ---------- + charge_spin + Frame conditions supplied by the caller, or ``None``. + nf + Number of frames on the flat node axis. + ref + Reference array supplying the compute namespace, dtype and device. + + Returns + ------- + Array or None + Frame conditions with shape ``(nf, 2)``, or ``None`` for a + descriptor without charge conditioning. + """ + if self.charge_spin_embedding is None: + return None + return canonicalize_charge_spin( + charge_spin, + self.default_chg_spin, + nf=nf, + ref=ref, + ) + @cast_precision def call( self, @@ -517,8 +595,9 @@ def call( Communication metadata accepted by the common descriptor ABI; unused. charge_spin - Charge/spin conditioning accepted by the common descriptor ABI; - unsupported and unused. + Frame-level total charge and spin multiplicity with shape + ``(F, 2)``. Read only when the descriptor is configured with + ``add_chg_spin_ebd``. Returns ------- @@ -539,7 +618,7 @@ def call( graph_from_dense_quartet, ) - del fparam, comm_dict, charge_spin + del fparam, comm_dict xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) nf, nloc, nnei = nlist.shape @@ -556,6 +635,8 @@ def call( graph, atype_local, self.type_embedding.call(), + None, + self.require_charge_spin(charge_spin, nf, graph.edge_vec), ) # === Step 3. Restore the common dense descriptor ABI === @@ -572,6 +653,7 @@ def evaluate_graph( atype: Array, type_embedding: Array, spin: Array | None = None, + charge_spin: Array | None = None, ) -> tuple[Array, Array]: """Evaluate the graph-native descriptor equations. @@ -589,6 +671,9 @@ def evaluate_graph( spin Per-node spin vectors with shape ``(N, 3)``, mandatory for a spin-conditioned descriptor and ignored otherwise. + charge_spin + Canonical frame conditions with shape ``(nf, 2)``, mandatory for + a charge-conditioned descriptor and ignored otherwise. Returns ------- @@ -602,6 +687,10 @@ def evaluate_graph( ValueError If the descriptor is spin conditioned and ``spin`` is absent. """ + from deepmd.dpmodel.utils.neighbor_graph import ( + frame_id_from_n_node, + ) + xp = array_api_compat.array_namespace(graph.edge_vec) # === Step 1. Place the precomputed type table in the graph namespace === @@ -617,10 +706,31 @@ def evaluate_graph( ) dst = graph.edge_index[1] n_total = atype.shape[0] - center_type_embedding = self.gather_rows(type_embedding, atype, xp) - pair_tables = self.pair_film.call(type_embedding) - # === Step 2. Condition the per-node spin === + # === Step 2. Embed the frame condition === + # The two emitted vectors reach the two places the type embedding + # enters: the centre type tail and the ordered pair encoder. Both + # frames and edges address them through the node-to-frame map, an + # edge inheriting the frame of the centre it reduces onto. + type_shift, pair_hidden_bias, edge_hidden_bias = None, None, None + if self.charge_spin_embedding is not None: + type_shift, pair_hidden_bias = self.charge_spin_embedding.call(charge_spin) + frame_index = frame_id_from_n_node(graph.n_node, n_total) + edge_hidden_bias = self.gather_rows( + pair_hidden_bias, + self.gather_rows(frame_index, dst, xp), + xp, + ) + type_shift = self.gather_rows(type_shift, frame_index, xp) + + center_type_embedding = self.build_center_type_features( + type_embedding, + atype, + type_shift, + ) + pair_latent = self.pair_film.pair_latent(type_embedding) + + # === Step 3. Condition the per-node spin === # The mask and the reference magnitude are applied once, so every # downstream spin route inherits them and the magnetic force of a # non-magnetic type vanishes identically rather than numerically. @@ -630,15 +740,16 @@ def evaluate_graph( else self.spin.conditioned_spin(self.require_spin(spin), atype) ) - # === Step 3. Build the masked edge amplitudes and harmonics === + # === Step 4. Build the masked edge amplitudes and harmonics === amplitude, basis, envelope, spin_payload = self.build_edge_features( graph, atype, - pair_tables, + pair_latent, + edge_hidden_bias, conditioned_spin, ) - # === Step 4. Reduce the degree-wise moments === + # === Step 5. Reduce the degree-wise moments === moments, divisors = self.aggregate_moments( amplitude, basis, @@ -651,7 +762,7 @@ def evaluate_graph( n_total, ) - # === Step 5. Build calibrated invariant features === + # === Step 6. Build calibrated invariant features === return ( self.build_invariant_descriptor( moments, @@ -661,11 +772,94 @@ def evaluate_graph( envelope[:, None], ) + def build_center_type_features( + self, + type_embedding: Array, + atype: Array, + type_shift: Array | None, + ) -> Array: + """Gather the centre type embedding and add the frame condition. + + The padding type keeps its zero row. Its output row is discarded, but + compressed inference conditions the real rows of a frozen type table + and leaves the padding row untouched, so shifting it here would break + the parity between the two paths. + + Parameters + ---------- + type_embedding + Complete type table with shape ``(ntypes + 1, channels)``. + atype + Flat node types with shape ``(N,)``. + type_shift + Per-node condition shift with shape ``(N, channels)``, or + ``None`` for a descriptor without charge conditioning. + + Returns + ------- + Array + Centre type features with shape ``(N, channels)``. + """ + xp = array_api_compat.array_namespace(type_embedding) + features = self.gather_rows(type_embedding, atype, xp) + if type_shift is None: + return features + real_type = xp.astype(atype < self.ntypes, features.dtype) + return features + type_shift * real_type[:, None] + + def build_pair_conditioning( + self, + pair_latent: tuple[Array, Array], + pair_index: Array, + edge_hidden_bias: Array | None, + ) -> tuple[Array, Array, Array | None, Array | None, Array | None]: + """Evaluate the ordered pair conditioning of every edge. + + The heads are applied on the coarsest axis over which their argument + is constant. Without a frame condition that axis is the ordered type + pair, so the finite cache of :math:`(T+1)^2` rows is built once and + gathered. With one, the argument additionally depends on the frame, + and the product axis is larger than the edge count for the molecular + systems a charge state describes, so the heads move to the edge axis. + Both routes evaluate the same function. + + Parameters + ---------- + pair_latent + Condition-independent ordered-pair state from + :meth:`~deepmd.dpmodel.descriptor.dpa4c_nn.pair_film.OrderedPairFiLM.pair_latent`. + pair_index + Ordered type-pair index of each edge with shape ``(E,)``. + edge_hidden_bias + Per-edge condition bias with shape ``(E, pair_hidden_width)``, or + ``None`` for a descriptor without charge conditioning. + + Returns + ------- + tuple + Radial scale and shift with shape ``(E, channels)``, the + mode-mixing matrices with shape ``(E, channels, radial_modes)``, + and the ordered spin scale and shift with shape + ``(E, spin_channels)``. The trailing three are ``None`` when + their mechanism is disabled. + """ + pre_activation, base_shift = pair_latent + if edge_hidden_bias is None: + return tuple( + None if table is None else self.gather_rows(table, pair_index) + for table in self.pair_film.heads(pre_activation, base_shift) + ) + return self.pair_film.heads( + self.gather_rows(pre_activation, pair_index) + edge_hidden_bias, + self.gather_rows(base_shift, pair_index), + ) + def build_edge_features( self, graph: Any, atype: Array, - pair_tables: tuple, + pair_latent: tuple[Array, Array], + edge_hidden_bias: Array | None = None, conditioned_spin: Array | None = None, ) -> tuple[Array, Array, Array, Array | None]: r"""Build the enveloped edge amplitudes, harmonics, and spin payload. @@ -692,14 +886,13 @@ def build_edge_features( Neighbor graph in descriptor compute precision. atype Flat node types with shape ``(N,)``. - pair_tables - Ordered pair cache produced by - :meth:`~deepmd.dpmodel.descriptor.dpa4c_nn.pair_film.OrderedPairFiLM.call`: - radial scale and shift with shape - ``((ntypes + 1) ** 2, channels)``, the mode-mixing table with - shape ``((ntypes + 1) ** 2, channels, radial_modes)`` or ``None``, - and the ordered spin scale and shift with shape - ``((ntypes + 1) ** 2, spin_channels)`` or ``None``. + pair_latent + Condition-independent ordered-pair state from + :meth:`~deepmd.dpmodel.descriptor.dpa4c_nn.pair_film.OrderedPairFiLM.pair_latent`. + edge_hidden_bias + Per-edge frame-condition bias with shape + ``(E, pair_hidden_width)``, or ``None`` for a descriptor without + charge conditioning. conditioned_spin Conditioned per-node spin with shape ``(N, 3)``, or ``None`` for a spin-free descriptor. @@ -720,8 +913,6 @@ def build_edge_features( apply_pair_exclusion, ) - pair_scale, pair_shift, pair_mixing, spin_scale, spin_shift = pair_tables - # === Step 1. Merge graph and descriptor-level exclusion masks === graph = apply_pair_exclusion(graph, atype, self.emask) xp = array_api_compat.array_namespace(graph.edge_vec) @@ -752,8 +943,11 @@ def build_edge_features( radial_hidden = self.radial_embedding.call_hidden(radial_basis) radial = self.radial_embedding.call_output(radial_hidden) pair_index = center_type * (self.ntypes + 1) + neighbor_type - scale = self.gather_rows(pair_scale, pair_index, xp) # (E, C) - shift = self.gather_rows(pair_shift, pair_index, xp) # (E, C) + scale, shift, mixing, spin_scale, spin_shift = self.build_pair_conditioning( + pair_latent, + pair_index, + edge_hidden_bias, + ) amplitude = radial * scale + shift # === Step 4. Add the pair-conditioned radial mode residual === @@ -763,8 +957,7 @@ def build_edge_features( # one GEMV per edge leaves the tiny C-by-R operands far short of # memory bandwidth, whereas the reduction is a plain streaming pass. # Expanding the ordered table per edge dominates the cost either way. - if pair_mixing is not None: - mixing = self.gather_rows(pair_mixing, pair_index, xp) # (E, C, R) + if mixing is not None: modes = self.radial_mode_head(radial_hidden) # (E, R) amplitude = amplitude + xp.sum(mixing * modes[:, None, :], axis=-1) @@ -784,7 +977,6 @@ def build_edge_features( envelope[:, 0], spin_scale, spin_shift, - pair_index, ) ) @@ -1058,6 +1250,7 @@ def share_params( "pair_film", "readout", "spin", + "charge_spin_embedding", ): setattr(self, name, getattr(base_class, name)) self.mean = base_class.mean @@ -1079,8 +1272,14 @@ def structure_signature(self) -> tuple: ``trainable`` decides whether those layers carry gradients at all; ``type_map`` fixes what the rows of the shared type table mean; ``use_spin`` fixes both the presence - and the row meaning of the shared spin tables. Precision itself enters - through its resolved dtype so that equivalent spellings agree. + and the row meaning of the shared spin tables; ``add_chg_spin_ebd`` + fixes the presence of the condition module and the width of the pair + encoder head it drives. Precision itself enters through its resolved + dtype so that equivalent spellings agree. + + ``default_chg_spin`` is deliberately absent: it is a fallback for a + missing input rather than a property of the shared parameters, so two + branches may legitimately default to different charge states. Branch-local state is deliberately absent. ``exclude_types`` is the only such field: it configures the pair-exclusion mask, which each @@ -1103,6 +1302,7 @@ def structure_signature(self) -> tuple: self.trainable, None if self.type_map is None else tuple(self.type_map), None if self.use_spin is None else tuple(self.use_spin), + self.add_chg_spin_ebd, np.dtype(PRECISION_DICT[self.precision]).name, ) @@ -1240,6 +1440,13 @@ def compute_input_stats( if frame["spin"] is None else xp.asarray(frame["spin"], dtype=dtype, device=device) ) + charge_spin = ( + None + if frame["charge_spin"] is None + else xp.asarray( + frame["charge_spin"], dtype=dtype, device=device + ) + ) graph = build_neighbor_graph( coord, atype, @@ -1250,6 +1457,7 @@ def compute_input_stats( graph, xp.reshape(atype, (-1,)), spin=None if spin is None else xp.reshape(spin, (-1, 3)), + charge_spin=charge_spin, ) output_np = to_numpy_array(output).reshape( -1, @@ -1349,15 +1557,17 @@ def _calibration_frames(self, system: dict) -> list[dict]: ---------- system Sampled system carrying ``coord``, ``atype``, an optional ``box``, - and, for a spin-conditioned descriptor, the per-atom moment under - either ``model_spin`` or ``spin``. + for a spin-conditioned descriptor the per-atom moment under either + ``model_spin`` or ``spin``, and for a charge-conditioned + descriptor the frame condition under ``charge_spin``. Returns ------- list[dict] One entry per drawn frame, each with a leading frame axis of - length one and a ``spin`` entry that is ``None`` for a spin-free - descriptor. + length one and ``spin`` and ``charge_spin`` entries that are + ``None`` when their mechanism is disabled or, for the frame + condition, when the descriptor falls back to its default. Raises ------ @@ -1383,6 +1593,21 @@ def _calibration_frames(self, system: dict) -> list[dict]: spin = ( None if spin is None else np.reshape(to_numpy_array(spin), (nframes, -1, 3)) ) + # The calibration must see the sampled distribution of charge states, + # because the fixed diagonal preconditioner it freezes has to hold for + # every one of them. A system without the key falls back to the + # configured default at the descriptor boundary. A system that states + # one condition for all of its frames is broadcast here, so that the + # calibration accepts exactly the shapes evaluation accepts. + charge_spin = None if not self.add_chg_spin_ebd else system.get("charge_spin") + charge_spin = ( + None + if charge_spin is None + else np.broadcast_to( + np.reshape(to_numpy_array(charge_spin), (-1, 2)), + (nframes, 2), + ) + ) indices = np.linspace( 0, nframes - 1, @@ -1395,6 +1620,9 @@ def _calibration_frames(self, system: dict) -> list[dict]: "atype": atype[index : index + 1], "box": None if box is None else box[index : index + 1], "spin": None if spin is None else spin[index : index + 1], + "charge_spin": ( + None if charge_spin is None else charge_spin[index : index + 1] + ), } for index in indices ] @@ -1470,8 +1698,15 @@ def serialize(self) -> dict: "type_map": self.type_map, "seed": self.seed, "use_spin": self.use_spin, + "add_chg_spin_ebd": self.add_chg_spin_ebd, + "default_chg_spin": self.default_chg_spin, "spin": None, "spin_channels": (None if self.spin is None else self.spin.serialize()), + "charge_spin_embedding": ( + None + if self.charge_spin_embedding is None + else self.charge_spin_embedding.serialize() + ), "type_embedding": self.type_embedding.serialize(), "radial_basis": self.radial_basis.serialize(), "radial_embedding": self.radial_embedding.serialize(), @@ -1523,6 +1758,7 @@ def deserialize(cls, data: dict) -> DescrptDPA4C: pair_film = data.pop("pair_film") readout = data.pop("readout") spin_channels = data.pop("spin_channels") + charge_spin_embedding = data.pop("charge_spin_embedding") obj = cls(**data) obj.type_embedding = SeZMTypeEmbedding.deserialize(type_embedding) @@ -1538,6 +1774,11 @@ def deserialize(cls, data: dict) -> DescrptDPA4C: obj.spin = ( None if spin_channels is None else SpinChannels.deserialize(spin_channels) ) + obj.charge_spin_embedding = ( + None + if charge_spin_embedding is None + else ChargeStateEmbedding.deserialize(charge_spin_embedding) + ) obj.set_stat_mean_and_stddev( variables["mean"], variables["stddev"], @@ -1689,6 +1930,29 @@ def supports_native_spin(self) -> bool: """Return whether ``call_graph`` conditions on a per-node spin.""" return self.spin is not None + def supports_charge_spin(self) -> bool: + """Return whether ``call_graph`` conditions on a frame charge state.""" + return self.charge_spin_embedding is not None + + def get_dim_chg_spin(self) -> int: + """Return the runtime width of the frame condition. + + Compression folds one charge state into the frozen type and ordered + pair tables, so the resulting snapshot evaluates that state and + consumes no runtime condition. Reporting zero is what routes such a + model onto the compact canonical lower, whose argument list carries + no conditioning slot. + """ + return 0 if self.compress or self.charge_spin_embedding is None else 2 + + def has_default_chg_spin(self) -> bool: + """Return whether a fallback frame condition is configured.""" + return self.default_chg_spin is not None + + def get_default_chg_spin(self) -> list[float] | None: + """Return the fallback ``[charge, multiplicity]``, if configured.""" + return self.default_chg_spin + def has_message_passing_across_ranks(self) -> bool: """Return whether intermediate halo communication is required.""" return False diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py index 75b73b4a9c..3d7bc17fd4 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py @@ -7,6 +7,11 @@ derive_bispectrum_ranks, enumerate_degree_triples, ) +from .charge_state import ( + ChargeStateEmbedding, + canonicalize_charge_spin, + validate_charge_state, +) from .geometry import ( MAX_ANGULAR_DEGREE, build_angular_basis, @@ -31,16 +36,19 @@ "MAX_ANGULAR_DEGREE", "NEIGHBOR_QUADRUPOLE_CHANNELS", "BispectrumLayout", + "ChargeStateEmbedding", "InvariantReadout", "OrderedPairFiLM", "SpinChannels", "build_angular_basis", "build_bispectrum_layout", "build_moment_indices", + "canonicalize_charge_spin", "degree_offsets", "derive_bispectrum_ranks", "derive_degree_channels", "derive_spin_channels", "enumerate_degree_triples", "packed_l2_to_stf", + "validate_charge_state", ] diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py b/deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py new file mode 100644 index 0000000000..79dbe52dd1 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py @@ -0,0 +1,374 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Frame-level charge and spin-multiplicity conditioning for DPA4C. + +The same nuclear geometry can be a cation, a neutral or an anion, and can be +a singlet or a triplet, with genuinely different energies and forces. The +descriptor therefore accepts one integer pair per frame, the total charge +:math:`Q_f` in units of the elementary charge and the spin multiplicity +:math:`M_f`, and turns it into two vectors that reach the two places where +the type embedding enters DPA4C: + +- :math:`\mathbf w_f`, added to the centre type embedding, which the + descriptor emits as its trailing output block; +- :math:`\mathbf y_f`, added to the hidden pre-activation of the ordered + type-pair encoder, which conditions the scale, shift, mode-mixing and spin + tables of every ordered pair. + +The second route is what makes the conditioning a property of the descriptor +rather than of the fitting network: it changes how a given geometry maps to +the degree-wise moments. It enters as a pre-activation bias rather than as a +shift of the encoder input because the first affine map of that encoder is +linear and bias free, so the two are equivalent while only the bias form +lets one shared projection over the finite type table serve every frame. + +Both routes are functions of the ordered type pair and of the condition, and +of nothing that depends on the interatomic distance. Compressed inference can +therefore fold them into its frozen type and ordered-pair tables without +changing a single table shape, at the cost of specializing the snapshot to +one charge state. + +This module is unrelated to :mod:`~deepmd.dpmodel.descriptor.dpa4c_nn.spin`, +which carries a per-atom magnetic moment. The two are independent inputs and +may be used together. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from ..dpa4_nn.embedding import ( + SeZMTypeEmbedding, +) +from ..dpa4_nn.mlp import ( + SwiGLUMLP, + resolve_swiglu_hidden_width, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + +#: Rows of the charge table, covering integer charges in ``[-100, 99]``. +CHARGE_TABLE_ROWS = 200 + +#: Index of the neutral charge row, so row ``CHARGE_OFFSET + Q`` holds ``Q``. +CHARGE_OFFSET = 100 + +#: Rows of the spin table, covering integer multiplicities below this bound. +MULTIPLICITY_TABLE_ROWS = 100 + +#: Half-open range of representable total charges, in units of the elementary +#: charge. +CHARGE_RANGE = (-CHARGE_OFFSET, CHARGE_TABLE_ROWS - CHARGE_OFFSET) + +#: Half-open range of representable spin multiplicities. +MULTIPLICITY_RANGE = (0, MULTIPLICITY_TABLE_ROWS) + + +def validate_charge_state(charge_spin: Any) -> list[float]: + """Check that a frame condition addresses a row of each embedding table. + + Both tables are indexed directly by the condition, and neither the gather + nor the compiled kernel bounds-checks that index, so an out-of-range value + would read past the table. Every host-side boundary that accepts a charge + state therefore passes it through here first. The per-forward path is + deliberately not guarded: its values come from the data pipeline, which + owns their validity exactly as it owns the validity of an atom type. + + Parameters + ---------- + charge_spin + A pair ``[charge, multiplicity]``, in any sequence form. + + Returns + ------- + list[float] + The same pair, as two floats. + + Raises + ------ + ValueError + If the pair does not hold exactly two integral values within the + representable ranges. + """ + values = [float(value) for value in np.reshape(np.asarray(charge_spin), (-1,))] + if len(values) != 2: + raise ValueError( + f"A charge state must be a `[charge, multiplicity]` pair, got " + f"{len(values)} values" + ) + for value, name, (low, high) in zip( + values, + ("charge", "multiplicity"), + (CHARGE_RANGE, MULTIPLICITY_RANGE), + strict=True, + ): + if value != int(value): + raise ValueError(f"The {name} must be an integer, got {value}") + if not low <= value < high: + raise ValueError( + f"The {name} must lie in [{low}, {high}), got {int(value)}" + ) + return values + + +class ChargeStateEmbedding(NativeOP): + r"""Embed the frame charge and spin multiplicity into two condition vectors. + + The two integers are embedded independently and mixed by one bias-free + SwiGLU trunk whose single output head is split into the two routes: + + .. math:: + + [\mathbf w_f \Vert \mathbf y_f] + = \operatorname{SwiGLU}\bigl( + [\,E^{Q}_{Q_f}\Vert E^{M}_{M_f}\,]W_{\rm in} + \bigr)W_{\rm out}. + + Charge and multiplicity share one nonlinear pathway rather than + contributing two additive embeddings. They are not independent degrees of + freedom: the number of unpaired electrons has the parity of the electron + count, so changing the charge by one flips it, and the structural + response to a spin-state change depends on the oxidation state it happens + in. An additive decomposition can represent neither coupling. + + The output projection is zero initialized, so an untrained descriptor is + independent of the condition for every value of it. This keeps the fixed + output calibration, which is measured once before training, free of a + random condition offset. + + Parameters + ---------- + channels + Width :math:`C_0` of the centre type embedding, and of each of the + two integer embedding tables. + pair_hidden_width + Width :math:`2H_{\rm pair}` of the ordered pair encoder's hidden + pre-activation. + precision + Parameter precision. + trainable + Whether the condition parameters receive optimizer updates. + seed + Random seed. + + Raises + ------ + ValueError + If ``channels`` or ``pair_hidden_width`` is not positive. + """ + + def __init__( + self, + channels: int, + pair_hidden_width: int, + *, + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: int | list[int] | None = None, + ) -> None: + if channels <= 0: + raise ValueError(f"`channels` must be positive, got {channels}") + if pair_hidden_width <= 0: + raise ValueError( + f"`pair_hidden_width` must be positive, got {pair_hidden_width}" + ) + self.channels = int(channels) + self.pair_hidden_width = int(pair_hidden_width) + self.precision = str(precision) + self.trainable = bool(trainable) + + self.charge_embedding = SeZMTypeEmbedding( + ntypes=CHARGE_TABLE_ROWS, + embed_dim=self.channels, + precision=self.precision, + seed=child_seed(seed, 0), + trainable=self.trainable, + padding=False, + ) + self.spin_embedding = SeZMTypeEmbedding( + ntypes=MULTIPLICITY_TABLE_ROWS, + embed_dim=self.channels, + precision=self.precision, + seed=child_seed(seed, 1), + trainable=self.trainable, + padding=False, + ) + hidden_width = resolve_swiglu_hidden_width(self.channels) + output_width = self.channels + self.pair_hidden_width + self.network = SwiGLUMLP( + [2 * self.channels, hidden_width, output_width], + precision=self.precision, + trainable=self.trainable, + seed=child_seed(seed, 2), + ) + # ``NativeLayer`` has no zero initializer; replicate it by overwriting + # the output projection, which leaves the condition inert until the + # optimizer moves it. + self.network.layers[-1].w = np.zeros( + (hidden_width, output_width), + dtype=PRECISION_DICT[self.precision.lower()], + ) + + def call(self, charge_spin: Array) -> tuple[Array, Array]: + """Embed the frame conditions into the two condition vectors. + + Parameters + ---------- + charge_spin + Frame conditions with shape ``(nf, 2)``, holding the total charge + and the spin multiplicity as exactly representable integers. + + Returns + ------- + type_shift + Centre type-embedding shift with shape ``(nf, channels)``. + pair_hidden_bias + Ordered pair encoder pre-activation bias with shape + ``(nf, pair_hidden_width)``. + """ + xp = array_api_compat.array_namespace(charge_spin) + charge = xp.astype(charge_spin[:, 0], xp.int64) + CHARGE_OFFSET + multiplicity = xp.astype(charge_spin[:, 1], xp.int64) + logits = self.network.call( + xp.concat( + (self.charge_embedding(charge), self.spin_embedding(multiplicity)), + axis=-1, + ) + ) + return logits[:, : self.channels], logits[:, self.channels :] + + def serialize(self) -> dict[str, Any]: + """Serialize the condition embedding. + + Returns + ------- + dict[str, Any] + Versioned configuration and the nested embedding tables and trunk. + """ + return { + "@class": "ChargeStateEmbedding", + "@version": 1, + "channels": self.channels, + "pair_hidden_width": self.pair_hidden_width, + "precision": self.precision, + "trainable": self.trainable, + "charge_embedding": self.charge_embedding.serialize(), + "spin_embedding": self.spin_embedding.serialize(), + "network": self.network.serialize(), + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> ChargeStateEmbedding: + """Deserialize a :class:`ChargeStateEmbedding`. + + Parameters + ---------- + data + Versioned dictionary produced by :meth:`serialize`. + + Returns + ------- + ChargeStateEmbedding + Reconstructed condition embedding. + + Raises + ------ + ValueError + If the payload does not describe a :class:`ChargeStateEmbedding`. + """ + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + if data.pop("@class") != "ChargeStateEmbedding": + raise ValueError("Invalid serialized class for ChargeStateEmbedding") + charge_embedding = data.pop("charge_embedding") + spin_embedding = data.pop("spin_embedding") + network = data.pop("network") + obj = cls(**data) + obj.charge_embedding = SeZMTypeEmbedding.deserialize(charge_embedding) + obj.spin_embedding = SeZMTypeEmbedding.deserialize(spin_embedding) + obj.network = SwiGLUMLP.deserialize(network) + return obj + + +def canonicalize_charge_spin( + charge_spin: Array | None, + default: list[float] | None, + *, + nf: int, + ref: Array, +) -> Array: + """Bring a frame-condition argument to the canonical ``(nf, 2)`` form. + + Parameters + ---------- + charge_spin + Frame conditions supplied by the caller, or ``None`` to fall back to + ``default``. A single pair is broadcast over the frame axis. + default + Configured fallback ``[charge, multiplicity]``, or ``None`` when the + descriptor requires an explicit condition. + nf + Number of frames. + ref + Reference array of the caller's compute context. The array namespace, + dtype and device are taken from it; deriving the namespace from the + NumPy ``default`` instead would break every non-NumPy backend. + + Returns + ------- + Array + Frame conditions with shape ``(nf, 2)``. + + Raises + ------ + ValueError + If no condition is available, or if the supplied condition does not + have shape ``(nf, 2)`` or a shape broadcastable to it. + """ + xp = array_api_compat.array_namespace(ref) + if charge_spin is None: + if default is None: + raise ValueError( + "A charge-conditioned DPA4C requires a frame `charge_spin`. " + "Set `default_chg_spin` to supply a fallback." + ) + charge_spin = xp.reshape( + xp.asarray( + np.asarray(default), + dtype=ref.dtype, + device=array_api_compat.device(ref), + ), + (1, 2), + ) + else: + charge_spin = xp.astype(xp.reshape(charge_spin, (-1, 2)), ref.dtype) + if charge_spin.shape[0] == 1 and nf != 1: + return xp.broadcast_to(charge_spin, (nf, 2)) + if charge_spin.shape[0] != nf: + raise ValueError( + f"`charge_spin` must hold one [charge, multiplicity] pair per " + f"frame, expected {nf} rows, got {charge_spin.shape[0]}" + ) + return charge_spin diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py b/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py index f8be2e5095..c9a28ae23b 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py @@ -155,10 +155,18 @@ def __init__( self.adam_spin_scale_anchor = (offset * signs[0]).astype(precision_dtype) self.adam_spin_shift_anchor = (offset * signs[1]).astype(precision_dtype) - def call( - self, type_embedding: Any - ) -> tuple[Any, Any, Any | None, Any | None, Any | None]: - """Build the ordered scale, shift, mixing, and spin tables. + @property + def pair_hidden_width(self) -> int: + """Return the width of the encoder's hidden pre-activation.""" + return 2 * self.hidden_dim + + def pair_latent(self, type_embedding: Any) -> tuple[Any, Any]: + """Build the condition-independent ordered-pair state. + + Both returned tables are functions of the ordered type pair alone. + Splitting them off lets a frame-level condition enter as an additive + pre-activation bias, which one shared projection over the finite type + table then serves for every frame. Parameters ---------- @@ -168,23 +176,14 @@ def call( Returns ------- - scale - Ordered radial scales with shape ``((T + 1) ** 2, channels)``. - shift - Ordered radial shifts with shape ``((T + 1) ** 2, channels)``. - mixing - Ordered mode-mixing matrices with shape - ``((T + 1) ** 2, channels, radial_modes)``, or ``None`` when - ``radial_modes`` is zero. - spin_scale - Ordered spin scales with shape - ``((T + 1) ** 2, spin_channels)``, or ``None`` when - ``spin_channels`` is zero. - spin_shift - Ordered spin shifts with the same shape as ``spin_scale``. + pre_activation + Hidden affine pre-activation with shape + ``((T + 1) ** 2, pair_hidden_width)``. + base_shift + Structural shift anchor :math:`T_a + T_b` with shape + ``((T + 1) ** 2, channels)``. """ xp = array_api_compat.array_namespace(type_embedding) - device = array_api_compat.device(type_embedding) ntypes = type_embedding.shape[0] pair_shape = (ntypes, ntypes, self.channels) pair_input = xp.reshape( @@ -197,17 +196,55 @@ def call( ), (-1, 2 * self.channels), ) - logits = self.network.call(pair_input) + return ( + self.network.call_hidden_affine(pair_input), + xp.reshape( + type_embedding[:, None, :] + type_embedding[None, :, :], + (-1, self.channels), + ), + ) + + def heads( + self, pre_activation: Any, base_shift: Any + ) -> tuple[Any, Any, Any | None, Any | None, Any | None]: + """Finish the conditioning tables from a hidden pre-activation. + + Every operation acts on the trailing axis, so the same head applies + to the finite ordered-pair table and to a per-edge expansion of it. + + Parameters + ---------- + pre_activation + Hidden affine pre-activation with shape + ``(..., pair_hidden_width)``. + base_shift + Structural shift anchor with shape ``(..., channels)``. + + Returns + ------- + scale + Radial scales with shape ``(..., channels)``. + shift + Radial shifts with shape ``(..., channels)``. + mixing + Mode-mixing matrices with shape + ``(..., channels, radial_modes)``, or ``None`` when + ``radial_modes`` is zero. + spin_scale + Spin scales with shape ``(..., spin_channels)``, or ``None`` when + ``spin_channels`` is zero. + spin_shift + Spin shifts with the same shape as ``spin_scale``. + """ + xp = array_api_compat.array_namespace(pre_activation) + device = array_api_compat.device(pre_activation) + logits = self.network.call_from_hidden_affine(pre_activation) # The output splits into the scale, the shift residual, the flattened # mixing matrix, and the two spin tables, in that order. shift_end = 2 * self.channels mixing_end = shift_end + self.channels * self.radial_modes spin_scale_end = mixing_end + self.spin_channels - base_shift = xp.reshape( - type_embedding[:, None, :] + type_embedding[None, :, :], - (-1, self.channels), - ) def anchored(anchor: Any, block: Any) -> Any: """Bound one spin head around its learned per-channel offset.""" @@ -218,29 +255,66 @@ def anchored(anchor: Any, block: Any) -> Any: anchor, dtype=block.dtype, device=device, - )[None, :] + ) ) return ( - 1.0 + xp.tanh(logits[:, : self.channels]), - base_shift + xp.tanh(logits[:, self.channels : shift_end]), + 1.0 + xp.tanh(logits[..., : self.channels]), + base_shift + xp.tanh(logits[..., self.channels : shift_end]), None if self.radial_modes == 0 else xp.reshape( - xp.tanh(logits[:, shift_end:mixing_end]), + xp.tanh(logits[..., shift_end:mixing_end]), (-1, self.channels, self.radial_modes), ), None if self.spin_channels == 0 else anchored( self.adam_spin_scale_anchor, - logits[:, mixing_end:spin_scale_end], + logits[..., mixing_end:spin_scale_end], ), None if self.spin_channels == 0 - else anchored(self.adam_spin_shift_anchor, logits[:, spin_scale_end:]), + else anchored(self.adam_spin_shift_anchor, logits[..., spin_scale_end:]), ) + def call( + self, type_embedding: Any, hidden_bias: Any = None + ) -> tuple[Any, Any, Any | None, Any | None, Any | None]: + """Build the ordered scale, shift, mixing, and spin tables. + + Parameters + ---------- + type_embedding + Complete type table with shape ``(T + 1, channels)``, where the + trailing row is the zero padding type. + hidden_bias + Optional frame-condition bias with shape + ``(pair_hidden_width,)``, added to the hidden pre-activation of + every ordered pair. ``None`` leaves the cache unconditioned. + + Returns + ------- + scale + Ordered radial scales with shape ``((T + 1) ** 2, channels)``. + shift + Ordered radial shifts with shape ``((T + 1) ** 2, channels)``. + mixing + Ordered mode-mixing matrices with shape + ``((T + 1) ** 2, channels, radial_modes)``, or ``None`` when + ``radial_modes`` is zero. + spin_scale + Ordered spin scales with shape + ``((T + 1) ** 2, spin_channels)``, or ``None`` when + ``spin_channels`` is zero. + spin_shift + Ordered spin shifts with the same shape as ``spin_scale``. + """ + pre_activation, base_shift = self.pair_latent(type_embedding) + if hidden_bias is not None: + pre_activation = pre_activation + hidden_bias + return self.heads(pre_activation, base_shift) + def serialize(self) -> dict[str, Any]: """Serialize the ordered pair encoder. diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py index ed936a1daf..d20075571c 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py @@ -469,7 +469,6 @@ def edge_payload( envelope: Array, pair_scale: Array, pair_shift: Array, - pair_index: Array, ) -> Array: r"""Build the per-edge spin payload reduced over neighbours. @@ -508,11 +507,9 @@ def edge_payload( envelope Masked C3 envelope with shape ``(E,)``. pair_scale - Ordered spin scales with shape ``((ntypes + 1) ** 2, spin_channels)``. + Per-edge ordered spin scales with shape ``(E, spin_channels)``. pair_shift - Ordered spin shifts with the same shape. - pair_index - Ordered type-pair index of each edge with shape ``(E,)``. + Per-edge ordered spin shifts with the same shape. Returns ------- @@ -523,11 +520,9 @@ def edge_payload( device = array_api_compat.device(conditioned_spin) channels = self.spin_channels neighbor_spin = xp.take(conditioned_spin, source, axis=0) # (E, 3) - scale = xp.take(pair_scale, pair_index, axis=0) # (E, Cs) - shift = xp.take(pair_shift, pair_index, axis=0) # (E, Cs) - spin_amplitude = (radial[:, :channels] * scale + shift) * (envelope * envelope)[ - :, None - ] + spin_amplitude = (radial[:, :channels] * pair_scale + pair_shift) * ( + envelope * envelope + )[:, None] neighbor_gate = xp.take( xp.take( xp_asarray_nodetach(xp, self.spin_mask, device=device), diff --git a/deepmd/kernels/cuda/dpa4c/graph_compress.py b/deepmd/kernels/cuda/dpa4c/graph_compress.py index bcb3d845b6..a551b2beea 100644 --- a/deepmd/kernels/cuda/dpa4c/graph_compress.py +++ b/deepmd/kernels/cuda/dpa4c/graph_compress.py @@ -37,6 +37,7 @@ dataclass, ) from typing import ( + TYPE_CHECKING, Any, ) @@ -50,9 +51,18 @@ derive_degree_channels, derive_spin_channels, packed_l2_to_stf, + validate_charge_state, ) +if TYPE_CHECKING: + from collections.abc import ( + Sequence, + ) + __all__ = [ + "CHARGE_STATE_ARTIFACTS", + "ChargeStateFold", + "build_charge_state_artifacts", "build_compression_artifacts", "build_radial_table", "coupling_records", @@ -71,6 +81,17 @@ SUPPORTED_LMAX = (2, 3, 4) SUPPORTED_RADIAL_MODES = (0, 2, 4, 8) +#: Compression artifacts that carry the frame charge state. They are the +#: complete image of a charge state in a compressed snapshot: overwriting +#: exactly these four re-specializes the snapshot to a different state, and +#: every other artifact depends on the trained weights alone. +CHARGE_STATE_ARTIFACTS = ( + "pair_film", + "pair_mixing", + "spin_pair", + "type_embedding", +) + # Degrees one and two carry the wide channel blocks and are contracted by # specialized closed forms in every backend; the remaining triples run through # the shared sparse coupling path. @@ -602,31 +623,6 @@ def build_compression_artifacts( table, info = build_radial_table(descriptor, stride) with torch.no_grad(): - type_embedding = descriptor.type_embedding.call().to( - device=device, - dtype=torch.float32, - ) - ( - pair_scale, - pair_shift, - pair_mixing, - spin_scale, - spin_shift, - ) = descriptor.pair_film.call(type_embedding) - pair_film = torch.stack((pair_scale, pair_shift), dim=-1) - spin_pair, spin_type = _build_spin_caches( - descriptor, - spin_scale, - spin_shift, - device, - ) - # The mode axis is innermost so that the coefficients a lane needs for - # one channel arrive in one or two vector loads. - mixing = ( - torch.zeros(0, dtype=torch.float32, device=device) - if pair_mixing is None - else pair_mixing.to(torch.float32) - ) readout_matrices = _build_readout_matrices(descriptor, profile, device) output_mean = descriptor.mean.to(device=device, dtype=torch.float32) output_inv_std = torch.reciprocal( @@ -641,11 +637,7 @@ def build_compression_artifacts( return { "data": table, "info": info, - "pair_film": pair_film.detach().contiguous(), - "pair_mixing": mixing.detach().contiguous(), - "spin_pair": spin_pair, - "spin_type": spin_type, - "type_embedding": type_embedding.detach().contiguous(), + "spin_type": _build_spin_type_cache(descriptor, device), "readout_matrices": readout_matrices, "coupling_meta": torch.as_tensor( _coupling_meta(records), @@ -664,65 +656,230 @@ def build_compression_artifacts( ), "output_mean": output_mean.detach().contiguous(), "output_inv_std": output_inv_std.detach().contiguous(), + **build_charge_state_artifacts(descriptor, descriptor.default_chg_spin), } -def _build_spin_caches( +def charge_state_artifacts( descriptor: Any, - spin_scale: torch.Tensor | None, - spin_shift: torch.Tensor | None, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor]: - """Freeze the two finite tables the native spin branch reads. + charge_spin: torch.Tensor | None, +) -> dict[str, torch.Tensor]: + """Build the compression artifacts that carry one frame charge state. + + The condition reaches the compiled kernel only through the ordered pair + encoder and the centre type table, neither of which depends on distance, + so these four artifacts are the complete image of a charge state in the + snapshot. Rebuilding them from a different state re-specializes the + snapshot without touching the radial table, the readout projections, the + angular couplings, the per-type spin scalars or the output calibration. + + This is the traceable form: it takes the condition as a tensor and reads + no host value, so the same code both builds the snapshot and, once + exported, rebuilds it on the deployment device. + :func:`build_charge_state_artifacts` is the host-side entry that validates + the condition first. + + Parameters + ---------- + descriptor + Evaluated pt_expt DPA4C descriptor. + charge_spin + Frame condition with shape ``(2,)``. Read only when the descriptor is + charge conditioned, and mandatory in that case. + + Returns + ------- + dict[str, torch.Tensor] + The artifacts named by :data:`CHARGE_STATE_ARTIFACTS`. + + Raises + ------ + ValueError + If a charge-conditioned descriptor is given no condition. + """ + device = next(descriptor.parameters()).device + type_embedding = descriptor.type_embedding.call().to( + device=device, + dtype=torch.float32, + ) + pair_hidden_bias, type_shift = None, None + if descriptor.charge_spin_embedding is not None: + if charge_spin is None: + raise ValueError( + "A charge-conditioned DPA4C snapshot must be built against a " + "frame condition. Set `default_chg_spin` to supply the state " + "the snapshot starts from." + ) + type_shift, pair_hidden_bias = descriptor.charge_spin_embedding.call( + charge_spin.to(dtype=torch.float32).reshape(1, 2) + ) + + # The ordered pair encoder reads the unconditioned type table and receives + # the condition as a pre-activation bias, while the centre type tail + # carries it as an additive shift on its real rows. Both mirror the + # portable path exactly. + ( + pair_scale, + pair_shift, + pair_mixing, + spin_scale, + spin_shift, + ) = descriptor.pair_film.call( + type_embedding, + hidden_bias=None if pair_hidden_bias is None else pair_hidden_bias[0], + ) + if type_shift is not None: + # The padding row keeps its zero embedding: the portable path masks + # the shift by atom type, so shifting it here would break the parity + # between the two routes on any system carrying padding or ghost + # nodes. + rows = torch.arange(type_embedding.shape[0], dtype=torch.int64, device=device) + real = rows < descriptor.ntypes + type_embedding = type_embedding + type_shift[0] * real[:, None] - ``spin_pair`` interleaves the ordered scale and shift so that one channel - arrives in a single 64-bit load, matching the geometric PairFiLM cache. - ``spin_type`` packs the four per-type scalars a node needs into one - 128-bit row: the gate divided by the reference magnitude, which conditions - the moment; the bare gate, which the magnetic-coordination family reads + empty = torch.zeros(0, dtype=torch.float32, device=device) + return { + "pair_film": torch.stack((pair_scale, pair_shift), dim=-1).contiguous(), + # The mode axis is innermost so that the coefficients a lane needs for + # one channel arrive in one or two vector loads. + "pair_mixing": ( + empty if pair_mixing is None else pair_mixing.to(torch.float32).contiguous() + ), + "spin_pair": ( + empty + if descriptor.spin is None + else torch.stack((spin_scale, spin_shift), dim=-1) + .to(device=device, dtype=torch.float32) + .contiguous() + ), + "type_embedding": type_embedding.contiguous(), + } + + +def build_charge_state_artifacts( + descriptor: Any, + charge_spin: Sequence[float] | None, +) -> dict[str, torch.Tensor]: + """Validate a frame charge state and build its compression artifacts. + + Parameters + ---------- + descriptor + Evaluated pt_expt DPA4C descriptor. + charge_spin + Frame condition ``[charge, multiplicity]``. Read only when the + descriptor is charge conditioned, and mandatory in that case. + + Returns + ------- + dict[str, torch.Tensor] + The artifacts named by :data:`CHARGE_STATE_ARTIFACTS`, detached. + + Raises + ------ + ValueError + If a charge-conditioned descriptor is given no condition, or one that + does not address a row of both embedding tables. + """ + state = None + if descriptor.charge_spin_embedding is not None and charge_spin is not None: + state = torch.tensor( + validate_charge_state(charge_spin), + dtype=torch.float32, + device=next(descriptor.parameters()).device, + ) + with torch.no_grad(): + return { + name: value.detach().contiguous() + for name, value in charge_state_artifacts(descriptor, state).items() + } + + +class ChargeStateFold(torch.nn.Module): + """Rebuild the charge-state artifacts of a compressed snapshot. + + Exporting this module beside the inference lower is what lets one + deployed artifact serve any charge state: the deployment layer runs it + once when the state becomes known and writes the four tensors over the + corresponding constants of the inference lower. The alternative, folding + inside the inference graph, would repeat an evaluation over + ``(T + 1) ** 2`` ordered pairs on every step for a value that is constant + over a molecular-dynamics run, and that evaluation does not shrink with + the system size. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor carrying charge conditioning. + """ + + def __init__(self, descriptor: Any) -> None: + super().__init__() + self.descriptor = descriptor + + def forward(self, charge_spin: torch.Tensor) -> tuple[torch.Tensor, ...]: + """Build the artifacts of one charge state. + + Parameters + ---------- + charge_spin + Frame condition with shape ``(1, 2)``, matching the tensor the + inference lower would receive. + + Returns + ------- + tuple[torch.Tensor, ...] + The artifacts named by :data:`CHARGE_STATE_ARTIFACTS`, in order. + """ + artifacts = charge_state_artifacts(self.descriptor, charge_spin.reshape(2)) + return tuple(artifacts[name] for name in CHARGE_STATE_ARTIFACTS) + + +def _build_spin_type_cache(descriptor: Any, device: torch.device) -> torch.Tensor: + """Freeze the per-type table of the native spin branch. + + The table packs the four per-type scalars a node needs into one 128-bit + row: the gate divided by the reference magnitude, which conditions the + moment; the bare gate, which the magnetic-coordination family reads because it counts neighbours that carry a moment rather than the moments - themselves; and the two on-site weights. + themselves; and the two on-site weights. None of them passes through the + ordered pair encoder, so unlike ``spin_pair`` this table is independent of + the frame charge state. Parameters ---------- descriptor Evaluated pt_expt DPA4C descriptor. - spin_scale, spin_shift - Ordered spin tables, or ``None`` for a spin-free descriptor. device - Device that receives the packed tables. + Device that receives the packed table. Returns ------- - spin_pair - Ordered cache with shape ``((T + 1) ** 2, spin_channels, 2)``, or an - empty tensor. - spin_type - Per-type table with shape ``(T + 1, 4)``, or an empty tensor. + torch.Tensor + Per-type table with shape ``(T + 1, 4)``, or an empty tensor for a + spin-free descriptor. """ - empty = torch.zeros(0, dtype=torch.float32, device=device) if descriptor.spin is None: - return empty, empty + return torch.zeros(0, dtype=torch.float32, device=device) spin = descriptor.spin gate = spin.spin_mask.to(device=device, dtype=torch.float32) reference = spin.spin_reference.to(device=device, dtype=torch.float32) - return ( - torch.stack((spin_scale, spin_shift), dim=-1) - .to(device=device, dtype=torch.float32) - .detach() - .contiguous(), - torch.stack( - ( - gate / reference, - gate, - spin.adam_spin_vector_weight.to(device=device, dtype=torch.float32), - spin.adam_spin_quadrupole_weight.to(device=device, dtype=torch.float32), - ), - dim=-1, + with torch.no_grad(): + return ( + torch.stack( + ( + gate / reference, + gate, + spin.adam_spin_vector_weight.to(device=device, dtype=torch.float32), + spin.adam_spin_quadrupole_weight.to( + device=device, dtype=torch.float32 + ), + ), + dim=-1, + ) + .detach() + .contiguous() ) - .detach() - .contiguous(), - ) def _build_readout_matrices( diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py index 1a9fd300b8..c94158808e 100644 --- a/deepmd/pt_expt/descriptor/dpa4c.py +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -133,6 +133,7 @@ def call_graph( type_embedding: torch.Tensor | None = None, comm_dict: dict | None = None, spin: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[torch.Tensor, None]: """Evaluate the graph descriptor with compressed CUDA dispatch. @@ -149,6 +150,11 @@ def call_graph( spin Per-node spin with shape ``(N, 3)``, mandatory for a spin-conditioned descriptor. + charge_spin + Frame-level charge state with shape ``(nf, 2)``. The compressed + branch does not read it: its frozen tables already carry the + charge state that compression baked in, which is why a compressed + descriptor reports a zero runtime condition width. Returns ------- @@ -189,6 +195,7 @@ def call_graph( type_embedding=type_embedding, comm_dict=comm_dict, spin=spin, + charge_spin=charge_spin, ) def build_edge_features( @@ -209,7 +216,10 @@ def build_edge_features( accumulates over the whole neighborhood and the readout raises the moments to the fourth power, so both stay in the descriptor compute precision; the ordered pair cache is likewise evaluated outside, over - the finite type table. + the finite type table. A charge-conditioned descriptor is the + exception: its conditioning heads run on the edge axis and therefore + inside the region, producing their bounded outputs in bfloat16 + alongside the amplitude they scale. Training follows ``use_amp`` and evaluation follows ``DP_AMP_INFER``. The two are independent: mixed precision at inference is a throughput @@ -429,6 +439,56 @@ def _set_compression( self.register_buffer(buffer_name, value) self.compress = True + def apply_charge_state(self, charge_spin: Any) -> None: + """Re-specialize a compressed snapshot to a frame charge state. + + The condition reaches the compiled kernel only through the ordered + pair encoder and the centre type table, so a charge state is fully + described by four of the frozen artifacts. Rebuilding those four + moves the snapshot to a different state at the cost of one evaluation + over the finite type table, leaving the radial table, the readout + projections, the angular couplings, the per-type spin scalars and the + output calibration untouched. + + The state is a constant of a molecular-dynamics run, so this is a + load-time operation. Applying it on every step would pay an + evaluation over ``(T + 1) ** 2`` ordered pairs for a value that never + changes, and that evaluation does not shrink with the system size. + + Parameters + ---------- + charge_spin + Frame condition ``[charge, multiplicity]``. + + Raises + ------ + RuntimeError + If the descriptor is not a compressed snapshot, or carries no + charge conditioning. + """ + if not self.compress: + raise RuntimeError( + "A charge state is applied to the frozen tables of a " + "compressed DPA4C snapshot; an uncompressed descriptor reads " + "the condition directly on every call." + ) + if self.charge_spin_embedding is None: + raise RuntimeError( + "This DPA4C was not built with `add_chg_spin_ebd`, so it has " + "no charge state to apply." + ) + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + build_charge_state_artifacts, + ) + + artifacts = build_charge_state_artifacts(self, charge_spin) + device = self.compress_pair_film.device + for name, value in artifacts.items(): + self._buffers[f"compress_{name}"] = value.to( + device=device, + dtype=torch.float32, + ).contiguous() + def set_stat_mean_and_stddev(self, mean: Any, stddev: Any) -> None: """Update output calibration and its compressed snapshot.""" super().set_stat_mean_and_stddev(mean, stddev) @@ -478,6 +538,11 @@ def enable_compression( ) -> None: """Build immutable artifacts for the current DPA4C mega kernel. + A charge-conditioned descriptor folds ``default_chg_spin`` into the + frozen type and ordered pair tables. The snapshot therefore evaluates + that one charge state at no runtime cost, and moves to another through + :meth:`apply_charge_state` rather than through a second compression. + Parameters ---------- min_nbor_dist diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 69f388b5da..90dfb17016 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -71,6 +71,9 @@ if TYPE_CHECKING: import ase.neighborlist + from torch._inductor.package import ( + AOTICompiledModel, + ) from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, @@ -135,6 +138,210 @@ def _reshape_charge_spin( ) from err +def _single_charge_state(charge_spin: np.ndarray, width: int) -> tuple[float, ...]: + """Reduce a requested condition to the one state a folded snapshot serves. + + A folded condition lives in tables that are shared by the whole snapshot, + so it is a property of the loaded model rather than of a frame. A request + that names one state per frame is honoured only when every frame names the + same one. + + Parameters + ---------- + charge_spin : np.ndarray + Requested condition, of any shape holding a whole number of states. + width : int + Number of values in one charge state. + + Returns + ------- + tuple[float, ...] + The requested state, with ``width`` values. + + Raises + ------ + ValueError + If the request does not hold at least one whole state, or holds + several states that are not all equal. + """ + values = np.asarray(charge_spin, dtype=np.float64) + if values.size == 0 or values.size % width: + raise ValueError( + f"charge_spin carries {values.size} values, which is not a positive " + f"whole number of {width}-wide charge states." + ) + states = values.reshape(-1, width) + if not bool((states == states[0]).all()): + raise ValueError( + "This model folds one charge state into its frozen tables and " + "therefore serves a single state at a time, but charge_spin names " + f"{states.shape[0]} states that are not all equal." + ) + return tuple(states[0].tolist()) + + +class _ChargeStateFold: + """The rebuild of the frozen tables that carry a compressed charge state. + + A compressed charge-conditioned descriptor evaluates its frame condition + once, when the model is frozen, into a handful of tables. Those tables + reach the compiled lower as module constants, so serving a different + condition means rebuilding them and writing them over those constants + rather than re-evaluating the condition on every step. The archive + therefore ships a second compiled artifact that performs the rebuild, + together with the name of the constant each of its outputs replaces. + + Every lower lifts its constants independently, so the names hold only for + the lower they were resolved against at freeze time. Only a compressed + DPA4C descriptor folds a charge state, and that family never carries + message passing across ranks, so an archive with a fold holds exactly one + lower and the question of a second set of names does not arise. + + Parameters + ---------- + model_file : str + Path to the ``.pt2`` archive. + metadata : dict[str, Any] + Parsed archive metadata. + target : AOTICompiledModel + The lower whose constants carry the condition. + + Attributes + ---------- + width : int + Number of values in a charge state this fold accepts. + + Raises + ------ + ValueError + If the archive declares a fold it cannot supply, or names no width + for a charge state. + """ + + def __init__( + self, + model_file: str, + metadata: dict[str, Any], + target: "AOTICompiledModel", + ) -> None: + import tempfile + import zipfile + + from torch._inductor import ( + aoti_load_package, + ) + + from deepmd.pt_expt.utils.serialization import ( + PT2_EXTRA_PREFIX, + ) + + self._constants = [str(name) for name in metadata["charge_state_constants"]] + # The lower reads no condition, so what the model accepts is the state + # the snapshot was frozen against, which is also the layout the rebuild + # consumes. + default_chg_spin = metadata.get("default_chg_spin") + if not default_chg_spin: + raise ValueError( + f"'{model_file}' ships a charge-state fold but names no " + "default_chg_spin, so the width of a charge state is unknown." + ) + self.width = len(default_chg_spin) + + entry = PT2_EXTRA_PREFIX + "charge_state.pt2" + with zipfile.ZipFile(model_file, "r") as zf: + if entry not in zf.namelist(): + raise ValueError( + f"Invalid .pt2 file '{model_file}': it declares " + f"charge_state_constants but carries no '{entry}', so it " + "cannot serve a runtime charge state." + ) + archive = zf.read(entry) + # ``aoti_load_package`` reads a path, so the nested archive is extracted + # to a temporary file that this object owns and releases with itself. + self._archive = tempfile.NamedTemporaryFile(suffix=".pt2") + self._archive.write(archive) + self._archive.flush() + self._runner = aoti_load_package(self._archive.name) + self._target = target + self._applied: tuple[float, ...] | None = None + + @classmethod + def load( + cls, + model_file: str, + metadata: dict[str, Any], + target: "AOTICompiledModel", + ) -> "_ChargeStateFold | None": + """Load the rebuild an archive declares, if it declares one. + + The constant-name field is the archive's claim that the rebuild ships + with it, so an archive that declares the names and cannot supply the + rebuild is malformed and fails here rather than degrading silently. + + Parameters + ---------- + model_file : str + Path to the ``.pt2`` archive. + metadata : dict[str, Any] + Parsed archive metadata. + target : AOTICompiledModel + The lower whose constants carry the condition. + + Returns + ------- + _ChargeStateFold or None + The fold, or ``None`` when the archive declares none. + """ + if "charge_state_constants" not in metadata: + return None + return cls(model_file, metadata, target) + + def apply(self, charge_spin: tuple[float, ...]) -> None: + """Rebuild the tables for a condition and write them over the constants. + + Rebuilding is skipped when the condition already applies, so a run that + evaluates many frames at one condition pays for it once. + + Applying a condition overwrites loaded module state and is therefore + not safe to interleave with a forward pass. + + Parameters + ---------- + charge_spin : tuple[float, ...] + The condition, with :attr:`width` values. + + Raises + ------ + RuntimeError + If the rebuild returns a different number of tables than the + archive names constants. + """ + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + if charge_spin == self._applied: + return + # The rebuild consumes the condition in the (1, width) float32 layout + # the inference lower would receive. + tables = self._runner( + torch.tensor([charge_spin], dtype=torch.float32, device=DEVICE) + ) + if len(tables) != len(self._constants): + raise RuntimeError( + f"The charge-state fold returned {len(tables)} tables but the " + f"archive names {len(self._constants)} constants; it cannot " + "serve a runtime charge state." + ) + # An unnamed output belongs to a mechanism this model has disabled and + # has no constant to reach. + self._target.load_constants( + {name: table for name, table in zip(self._constants, tables) if name}, + check_full_update=False, + ) + self._applied = charge_spin + + def _warn_legacy_edge_vec(metadata: dict) -> None: """Warn once per model load when an edge_vec-schema artifact is opened. @@ -223,6 +430,9 @@ def __init__( # identifies the lower ABI. self._neighbor_graph_method = neighbor_graph_method self._is_pt2 = model_file.endswith(".pt2") + # Only a compressed ``.pt2`` folds its charge state into constants; a + # model that reads the condition as an ordinary input needs no rebuild. + self._charge_state_fold: _ChargeStateFold | None = None if self._is_pt2: self._load_pt2(model_file) @@ -501,7 +711,9 @@ def _load_pt2(self, model_file: str) -> None: Archive entries are located under ``model/extra/`` so that the PyTorch 2.11 ``load_pt2`` loader accepts the archive without the - "outdated pt2 file" fallback warning. + "outdated pt2 file" fallback warning. A compressed charge-conditioned + archive carries a second compiled artifact beside the inference lower, + which :class:`_ChargeStateFold` loads to serve a runtime condition. """ import zipfile @@ -544,6 +756,10 @@ def _load_pt2(self, model_file: str) -> None: self._pt2_runner = aoti_load_package(model_file) self.exported_module = None + self._charge_state_fold = _ChargeStateFold.load( + model_file, self.metadata, self._pt2_runner + ) + def _load_pt(self, model_file: str, head: str | None = None) -> None: """Load a `.pt` training checkpoint (eager mode, no torch.export).""" from copy import ( @@ -922,19 +1138,116 @@ def has_default_chg_spin(self) -> bool: ) def get_dim_chg_spin(self) -> int: - """Get the width of charge/spin condition inputs.""" + """Get the width of the conditioning input of the compiled forward. + + This gates whether a forward pass is handed a condition tensor, and is + zero both for a model that carries no charge/spin conditioning and for + a compressed one, whose condition lives in frozen tables rather than in + an input. + """ if self._dpmodel is not None: return self._dpmodel.get_dim_chg_spin() return int(self.metadata.get("dim_chg_spin", 0)) + def _no_runtime_condition_reason(self) -> str: + """Explain why the loaded model serves no runtime charge state. + + A conditioning width of zero has two causes that call for different + answers. A model built without a charge state embedding carries no + condition at all. A charge-conditioned one reports zero because + compression folded its state into frozen tables: it still carries a + condition, but moving to another one means rebuilding those tables, + and only a ``.pt2`` archive ships that rebuild beside its inference + lower. + + Returns + ------- + str + The reason, phrased to complete a sentence about this model. + """ + if not self.has_chg_spin_ebd(): + return "this model carries no charge/spin conditioning." + return ( + "this model's charge state is folded into its compressed tables " + "rather than read as an input, and the artifact it was loaded " + "from ships no rebuild of those tables, so it serves only the " + "state it was compressed against. Freeze the compressed model as " + "a .pt2 archive, which carries that rebuild." + ) + + def _apply_charge_state(self, charge_spin: np.ndarray | None = None) -> None: + """Serve a charge/spin condition that the compiled forward cannot read. + + A compressed descriptor folds its condition into frozen tables that + reach the lower as constants, leaving the forward with no conditioning + argument. Rebuilding those tables is the whole mechanism for such a + model, and it travels with the archive rather than with the loader, so + a requested condition that no rebuild can reach is an error rather + than an argument to drop. + + Rebuilding overwrites loaded module state and is therefore not safe to + interleave with a forward pass; evaluation is single-threaded. + + Parameters + ---------- + charge_spin : np.ndarray, optional + The requested condition. A folded model falls back to the state + its snapshot was frozen against when none is given. + + Raises + ------ + ValueError + If the loaded model cannot serve the requested condition. + """ + fold = self._charge_state_fold + if fold is None: + if charge_spin is not None and self.get_dim_chg_spin() == 0: + raise ValueError( + f"charge_spin was given, but {self._no_runtime_condition_reason()}" + ) + return + if charge_spin is None: + charge_spin = self.metadata["default_chg_spin"] + fold.apply(_single_charge_state(charge_spin, fold.width)) + def _make_charge_spin_input( self, nframes: int, charge_spin: np.ndarray | None = None ) -> torch.Tensor | None: - """Build the fixed charge/spin tensor used by exported SeZM models.""" + """Serve a charge/spin condition and build the input the forward reads. + + The condition reaches the model by one of two routes, and this takes + whichever the model was frozen with, so that a caller does not depend + on it. A forward that reads the condition as an ordinary input gets it + as the returned tensor; a forward compiled without that input has the + condition applied to its constants by :meth:`_apply_charge_state` and + receives no tensor. + + Parameters + ---------- + nframes : int + Number of frames the returned tensor covers. + charge_spin : np.ndarray, optional + The requested condition, reshape-compatible with + ``(nframes, dim_chg_spin)``. Defaults to the condition stored in + the model. + + Returns + ------- + torch.Tensor or None + The condition with shape ``(nframes, dim_chg_spin)``, or ``None`` + when the compiled forward takes no conditioning input. + + Raises + ------ + ValueError + If the model reads a condition as an input and neither the caller + nor the model supplies one. + """ from deepmd.pt_expt.utils.env import ( DEVICE, ) + self._apply_charge_state(charge_spin) dim_chg_spin = self.get_dim_chg_spin() if dim_chg_spin == 0: return None @@ -2117,6 +2430,9 @@ def _eval_model_graph_spin( "charge_spin inputs; a model requiring them must not be " "frozen with a canonical lower kind." ) + # A compressed descriptor still serves a condition, through the + # constants its fold rewrites rather than through an argument. + self._apply_charge_state(charge_spin) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) @@ -2284,6 +2600,9 @@ def _eval_model_graph( "frozen with a canonical lower kind (the export eligibility " "gate should have rejected it)." ) + # A compressed descriptor still serves a condition, through the + # constants its fold rewrites rather than through an argument. + self._apply_charge_state(charge_spin) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 94250f36fa..3c896fd0d1 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -2043,6 +2043,189 @@ def _deserialize_to_file_pte( torch.export.save(exported, model_file, extra_files=extra_files) +def _charge_state_descriptor(data: dict, metadata: dict) -> Any | None: + """Return the descriptor whose charge state a deployment can rebuild. + + Only a compressed charge-conditioned descriptor qualifies. An + uncompressed one reads its frame condition as an ordinary input, so + serving another condition costs nothing and needs no rebuild. + + That combination is already visible in the collected metadata: a model + that carries a charge state embedding and yet reports a conditioning + width of zero has folded the condition into frozen tables. Every other + export answers from those two fields, and only a candidate pays for + rebuilding the model to reach its descriptor. + + Parameters + ---------- + data + Serialized model dictionary, as passed to the export entry point. + This is the dictionary the exported lower was traced from, whose + descriptor therefore owns the constants a fold would rewrite. + metadata + Archive metadata collected from that same model. + + Returns + ------- + Any or None + The evaluated descriptor, or ``None`` when the model carries no + compressed charge conditioning. + """ + if not metadata["has_chg_spin_ebd"] or metadata["dim_chg_spin"] != 0: + return None + + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + model = BaseModel.deserialize(data["model"]) + descriptor = getattr(getattr(model, "atomic_model", None), "descriptor", None) + if ( + descriptor is None + or not getattr(descriptor, "compress", False) + or getattr(descriptor, "charge_spin_embedding", None) is None + ): + return None + model.to("cpu") + model.eval() + return descriptor + + +def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: + """Name, per fold output, the constant of ``exported`` it replaces. + + A compressed charge-conditioned descriptor carries its frame condition in + four of its frozen tables, which reach a compiled lower as lifted + constants. A deployment can therefore serve any charge state by rebuilding + those four once, when the state becomes known, and writing them over the + corresponding constants. + + The result is positional: entry ``i`` names the constant that output ``i`` + of the fold replaces, so a consumer writes the outputs back without + knowing what any of them mean. An empty entry marks an output that no + constant receives, which is how a disabled mechanism appears: it + contributes an empty artifact and carries no charge state. + + ``make_fx`` traces a plain function, so it names every lifted tensor + positionally and the buffer names are gone by the time the program is + exported. The names are therefore recovered by value against the + descriptor's own artifacts, and an artifact matching anything other than + exactly one constant is an error rather than a guess. + + Each exported lower lifts its constants independently, so the names hold + only for the lower they were resolved against. Only a compressed DPA4C + descriptor folds a charge state, and that family never carries message + passing across ranks, so an archive with a fold holds exactly one lower. + + Parameters + ---------- + descriptor + Compressed charge-conditioned descriptor, as returned by + :func:`_charge_state_descriptor`. + exported + Exported lower whose constants are to be named. + + Returns + ------- + list[str] + Per fold output, the name of the constant it replaces. + + Raises + ------ + RuntimeError + If an artifact does not match exactly one lifted constant. + """ + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + CHARGE_STATE_ARTIFACTS, + ) + + lifted = { + name: tensor.detach().cpu() + for name, tensor in exported.state_dict.items() + if isinstance(tensor, torch.Tensor) + } + constants: list[str] = [] + for name in CHARGE_STATE_ARTIFACTS: + artifact = getattr(descriptor, f"compress_{name}").detach().cpu() + if artifact.numel() == 0: + constants.append("") + continue + matches = [ + key + for key, value in lifted.items() + if value.shape == artifact.shape + and value.dtype == artifact.dtype + and torch.equal(value, artifact) + ] + if len(matches) != 1: + raise RuntimeError( + f"The compressed artifact {name!r} matches {len(matches)} " + "constants of the exported lower; a runtime charge state " + "needs exactly one so that the deployment knows which " + "constant to overwrite." + ) + constants.append(matches[0]) + return constants + + +def _compile_charge_state_fold( + descriptor: Any, + aoti_configs: dict, +) -> bytes: + """Compile the rebuild of the charge-state artifacts as its own archive. + + Compiling the rebuild beside the inference lower is what lets one + deployed artifact serve any charge state: the deployment runs it once when + the state becomes known, then writes its outputs over the constants named + by :func:`_match_charge_state_constants`. + + Parameters + ---------- + descriptor + Compressed charge-conditioned descriptor, as returned by + :func:`_charge_state_descriptor`. + aoti_configs + Inductor options the inference lower was compiled with. + + Returns + ------- + bytes + The compiled fold archive. + """ + import os + import tempfile + + from torch._inductor import ( + aoti_compile_and_package, + ) + + import deepmd.pt_expt.utils.env as _env + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + ChargeStateFold, + ) + + log.info("Compiling the charge-state fold...") + # The descriptor is evaluated on the host, so the fold traces there and is + # moved to the target device with the rest of the program below. + sample = torch.tensor( + [descriptor.get_default_chg_spin()], + dtype=torch.float32, + device="cpu", + ) + fold = torch.export.export(ChargeStateFold(descriptor), (sample,)) + if _env.DEVICE.type != "cpu": + from torch.export.passes import ( + move_to_device_pass, + ) + + fold = move_to_device_pass(fold, _env.DEVICE) + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "charge_state.pt2") + aoti_compile_and_package(fold, package_path=path, inductor_configs=aoti_configs) + with open(path, "rb") as archive: + return archive.read() + + def _deserialize_to_file_pt2( model_file: str, data: dict, @@ -2145,6 +2328,19 @@ def _deserialize_to_file_pt2( exported, package_path=model_file, inductor_configs=aoti_configs ) + # Charge-state fold. Present only for a compressed charge-conditioned + # descriptor, whose frozen tables the deployment rebuilds once when the + # state becomes known. + charge_state_descriptor = _charge_state_descriptor(data, metadata) + charge_state_bytes: bytes | None = None + if charge_state_descriptor is not None: + metadata["charge_state_constants"] = _match_charge_state_constants( + charge_state_descriptor, exported + ) + charge_state_bytes = _compile_charge_state_fold( + charge_state_descriptor, aoti_configs + ) + # Second artifact: with-comm. Only for descriptors whose message # passing extends across rank boundaries. The flag was computed # from the model in ``_collect_metadata`` and is already in @@ -2202,3 +2398,5 @@ def _deserialize_to_file_pt2( zf.writestr( PT2_EXTRA_PREFIX + "forward_lower_with_comm.pt2", with_comm_bytes ) + if charge_state_bytes is not None: + zf.writestr(PT2_EXTRA_PREFIX + "charge_state.pt2", charge_state_bytes) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 4ff14d03b1..14b768320a 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -557,6 +557,34 @@ def descrpt_dpa4c_args() -> list[Argument]: default=True, doc="Whether descriptor parameters are trainable.", ), + Argument( + "add_chg_spin_ebd", + bool, + optional=True, + default=False, + doc=( + "Whether to condition the descriptor on the frame-level " + "`charge_spin` input `[charge, multiplicity]` of shape " + "`[nframes, 2]`. The embedded condition is added to the " + "center type embedding and to the hidden state of the " + "ordered type-pair encoder, so it changes how a given " + "geometry maps to the degree-wise moments. This is unrelated " + "to `model.spin`, which carries a per-atom magnetic moment." + ), + ), + Argument( + "default_chg_spin", + list[float], + optional=True, + default=None, + doc=( + "Fallback `[charge, multiplicity]` used when `charge_spin` is " + "absent from the input data. Only read when " + "`add_chg_spin_ebd` is enabled. Compression folds this value " + "into the frozen tables, so a compressed model evaluates " + "exactly this charge state and requires the option to be set." + ), + ), Argument( "seed", [int, None], diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 092cb30730..78f8c6f1e0 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -13,7 +13,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 29 +#define DP_C_API_VERSION 30 /** * @brief Neighbor list. @@ -2612,6 +2612,20 @@ int DP_DeepPotGetDimAParam(DP_DeepPot* dp); */ int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp); +/** + * @brief Fix the charge/spin condition served for the rest of the run. + * It becomes the condition of every later evaluation that is not given one + * explicitly. Intended to be called once, before the first evaluation. + * @param[in] dp The DP to use. + * @param[in] charge_spin The condition. + * @param[in] numb_chg_spin The number of values in charge_spin; must equal + * DP_DeepPotGetDimChgSpin. + * @since API version 30 + */ +extern void DP_DeepPotSetChargeSpin(DP_DeepPot* dp, + const double* charge_spin, + const int numb_chg_spin); + /** * @brief Check whether the atomic dimension of atomic parameters is nall * instead of nloc. @@ -2667,6 +2681,20 @@ int DP_DeepPotModelDeviGetDimAParam(DP_DeepPotModelDevi* dp); */ int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp); +/** + * @brief Fix the charge/spin condition served for the rest of the run. + * Applied to every model, so that the deviation is taken between models under + * the same condition. + * @param[in] dp The DP to use. + * @param[in] charge_spin The condition. + * @param[in] numb_chg_spin The number of values in charge_spin; must equal + * DP_DeepPotModelDeviGetDimChgSpin. + * @since API version 30 + */ +extern void DP_DeepPotModelDeviSetChargeSpin(DP_DeepPotModelDevi* dp, + const double* charge_spin, + const int numb_chg_spin); + /** * @brief Check whether the atomic dimension of atomic parameters is nall * instead of nloc. @@ -2765,6 +2793,20 @@ int DP_DeepSpinGetDimAParam(DP_DeepSpin* dp); */ int DP_DeepSpinGetDimChgSpin(DP_DeepSpin* dp); +/** + * @brief Fix the charge/spin condition served for the rest of the run. + * It becomes the condition of every later evaluation that is not given one + * explicitly. Intended to be called once, before the first evaluation. + * @param[in] dp The DP Spin Model to use. + * @param[in] charge_spin The condition. + * @param[in] numb_chg_spin The number of values in charge_spin; must equal + * DP_DeepSpinGetDimChgSpin. + * @since API version 30 + */ +extern void DP_DeepSpinSetChargeSpin(DP_DeepSpin* dp, + const double* charge_spin, + const int numb_chg_spin); + /** * @brief Check whether the atomic dimension of atomic parameters is nall * instead of nloc. @@ -2826,6 +2868,20 @@ int DP_DeepSpinModelDeviGetDimAParam(DP_DeepSpinModelDevi* dp); */ int DP_DeepSpinModelDeviGetDimChgSpin(DP_DeepSpinModelDevi* dp); +/** + * @brief Fix the charge/spin condition served for the rest of the run. + * Applied to every model, so that the deviation is taken between models under + * the same condition. + * @param[in] dp The DP Spin Model Deviation to use. + * @param[in] charge_spin The condition. + * @param[in] numb_chg_spin The number of values in charge_spin; must equal + * DP_DeepSpinModelDeviGetDimChgSpin. + * @since API version 30 + */ +extern void DP_DeepSpinModelDeviSetChargeSpin(DP_DeepSpinModelDevi* dp, + const double* charge_spin, + const int numb_chg_spin); + /** * @brief Check whether the atomic dimension of atomic parameters is nall * instead of nloc. diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index e88a38d93a..681b197b0e 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1287,6 +1287,19 @@ class DeepPot : public DeepBaseModel { return dchgspin; } + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * It becomes the condition of every later evaluation that is not given one + * explicitly. Intended to be called once, before the first evaluation. + * @param[in] charge_spin The condition, of length dim_chg_spin(). + **/ + void set_charge_spin(const std::vector& charge_spin) { + assert(dp); + DP_DeepPotSetChargeSpin(dp, charge_spin.data(), + static_cast(charge_spin.size())); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + } + /** * @brief Evaluate a device-resident edge graph with FP64 edge vectors. * @@ -1889,6 +1902,19 @@ class DeepSpin : public DeepBaseModel { return dchgspin; } + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * It becomes the condition of every later evaluation that is not given one + * explicitly. Intended to be called once, before the first evaluation. + * @param[in] charge_spin The condition, of length dim_chg_spin(). + **/ + void set_charge_spin(const std::vector& charge_spin) { + assert(dp); + DP_DeepSpinSetChargeSpin(dp, charge_spin.data(), + static_cast(charge_spin.size())); + DP_CHECK_OK(DP_DeepSpinCheckOK, dp); + } + /** * @brief Evaluate a compact canonical graph on the model device. * @@ -2582,6 +2608,19 @@ class DeepPotModelDevi : public DeepBaseModelDevi { return dchgspin; } + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * Applied to every model, so that the deviation is taken between models + * under the same condition. + * @param[in] charge_spin The condition, of length dim_chg_spin(). + **/ + void set_charge_spin(const std::vector& charge_spin) { + assert(dp); + DP_DeepPotModelDeviSetChargeSpin(dp, charge_spin.data(), + static_cast(charge_spin.size())); + DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); + } + /** * @brief Evaluate the energy, force and virial by using this DP model *deviation. @@ -3065,6 +3104,19 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { return dchgspin; } + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * Applied to every model, so that the deviation is taken between models + * under the same condition. + * @param[in] charge_spin The condition, of length dim_chg_spin(). + **/ + void set_charge_spin(const std::vector& charge_spin) { + assert(dp); + DP_DeepSpinModelDeviSetChargeSpin(dp, charge_spin.data(), + static_cast(charge_spin.size())); + DP_CHECK_OK(DP_DeepSpinModelDeviCheckOK, dp); + } + /** * @brief Evaluate the energy, force, magnetic force and virial by using this *DP spin model deviation. diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index cfb2ed5b9e..cee85e258b 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -771,6 +771,26 @@ bool validate_model_devi_nframes(DP_DeepBaseModelDevi* dp, const int nframes) { return false; } +/** + * @brief Copy a charge/spin condition supplied by a C caller. + * + * The C boundary carries a bare pointer, so the element count the caller + * passes alongside it is the only thing standing between a short array and a + * read past its end, or a long one and a silent truncation. Callers invoke + * this inside DP_REQUIRES_OK so that a mismatch is reported through + * DP_*CheckOK rather than unwinding into C. + */ +std::vector copy_charge_spin(const double* charge_spin, + const int numb_chg_spin, + const int dim_chg_spin) { + if (numb_chg_spin != dim_chg_spin) { + throw deepmd::deepmd_exception( + "the charge/spin condition carries " + std::to_string(numb_chg_spin) + + " values but the model expects " + std::to_string(dim_chg_spin)); + } + return std::vector(charge_spin, charge_spin + numb_chg_spin); +} + } // namespace template @@ -2757,8 +2777,22 @@ int DP_DeepPotGetDimAParam(DP_DeepPot* dp) { int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp) { return dp->dp.dim_chg_spin(); } +void DP_DeepPotSetChargeSpin(DP_DeepPot* dp, + const double* charge_spin, + const int numb_chg_spin) { + DP_REQUIRES_OK(dp, dp->dp.set_charge_spin(copy_charge_spin( + charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); +} + int DP_DeepSpinGetDimChgSpin(DP_DeepSpin* dp) { return dp->dp.dim_chg_spin(); } +void DP_DeepSpinSetChargeSpin(DP_DeepSpin* dp, + const double* charge_spin, + const int numb_chg_spin) { + DP_REQUIRES_OK(dp, dp->dp.set_charge_spin(copy_charge_spin( + charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); +} + bool DP_DeepPotIsAParamNAll(DP_DeepPot* dp) { return DP_DeepBaseModelIsAParamNAll(static_cast(dp)); } @@ -2799,6 +2833,13 @@ int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp) { return dp->dp.dim_chg_spin(); } +void DP_DeepPotModelDeviSetChargeSpin(DP_DeepPotModelDevi* dp, + const double* charge_spin, + const int numb_chg_spin) { + DP_REQUIRES_OK(dp, dp->dp.set_charge_spin(copy_charge_spin( + charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); +} + bool DP_DeepPotModelDeviIsAParamNAll(DP_DeepPotModelDevi* dp) { return DP_DeepBaseModelDeviIsAParamNAll( static_cast(dp)); @@ -2878,6 +2919,13 @@ int DP_DeepSpinModelDeviGetDimChgSpin(DP_DeepSpinModelDevi* dp) { return dp->dp.dim_chg_spin(); } +void DP_DeepSpinModelDeviSetChargeSpin(DP_DeepSpinModelDevi* dp, + const double* charge_spin, + const int numb_chg_spin) { + DP_REQUIRES_OK(dp, dp->dp.set_charge_spin(copy_charge_spin( + charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); +} + bool DP_DeepSpinModelDeviIsAParamNAll(DP_DeepSpinModelDevi* dp) { return DP_DeepBaseModelDeviIsAParamNAll( static_cast(dp)); diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 8b2d7ae994..56cbef8704 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -206,6 +206,27 @@ class DeepPotBackend : public DeepBaseModelBackend { **/ virtual int dim_chg_spin() const { return 0; } + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * + * An override is needed only where the condition has to be folded into state + * that is built ahead of the evaluations using it, as in a compressed model + * whose tables are specialized to one condition. A backend that reads the + * condition as an ordinary per-call input has nothing to install, so the + * default is a no-op rather than an error: the condition still reaches such + * a backend on every evaluation, through the charge_spin argument of + * computew(). The request is refused only by a model that carries no + * charge/spin conditioning at all, which no route can honour. + * + * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. + **/ + virtual void set_charge_spin(const std::vector& charge_spin) { + if (dim_chg_spin() == 0) { + throw deepmd::deepmd_exception( + "this model does not support a charge/spin condition"); + } + } + // charge_spin-aware computew overloads. Default implementations call the // existing pure-virtual overloads (ignoring charge_spin) so that backends // that do not support charge/spin do not need any changes. DeepPotPTExpt @@ -841,6 +862,12 @@ class DeepPot : public DeepBaseModel { int dim_chg_spin() const; + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. + **/ + void set_charge_spin(const std::vector& charge_spin); + protected: std::shared_ptr dp; }; @@ -885,6 +912,18 @@ class DeepPotModelDevi : public DeepBaseModelDevi { return numb_models > 0 ? dps[0]->dim_chg_spin() : 0; }; + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * Applied to every model, so that the deviation is taken between models + * under the same condition. + * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. + **/ + void set_charge_spin(const std::vector& charge_spin) { + for (unsigned ii = 0; ii < dps.size(); ++ii) { + dps[ii]->set_charge_spin(charge_spin); + } + }; + /** * @brief Evaluate the energy, force and virial by using these DP models. * @param[out] all_ener The system energies of all models. diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index cde5b68cd9..361741376b 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -16,11 +16,12 @@ #include "DeepPot.h" -// Forward-declare to keep TempFile out of public header. Defined in +// Forward-declare to keep these out of the public header. Defined in // commonPTExpt.h. namespace deepmd::ptexpt { class TempFile; -} +class ChargeStateFold; +} // namespace deepmd::ptexpt namespace torch::inductor { class AOTIModelPackageLoader; @@ -117,10 +118,37 @@ class DeepPotPTExpt : public DeepPotBackend { assert(inited); return daparam; }; + /** + * @brief The width of a charge/spin condition this model accepts. + * + * This is the width a caller names a condition with, which is not in + * general the width of the conditioning input of the compiled forward: + * compression folds the condition into frozen tables and so removes it from + * the argument list, leaving a model that still serves a condition through + * the fold shipped beside the inference lower. + **/ int dim_chg_spin() const override { assert(inited); - return dchgspin; + return settable_chgspin; }; + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * + * The condition reaches the model by one of two routes, and this sets both + * so that the caller does not depend on how the model was frozen. It + * becomes the condition of every later forward pass that is not given one + * explicitly. A compressed descriptor additionally carries the condition + * inside frozen tables that the compiled lower holds as constants; when the + * archive ships the fold that rebuilds them, it runs here and the resulting + * tables are written over those constants. + * + * Intended to be called once, before inference, since overwriting the + * constants of a loaded module is not safe to interleave with a forward + * pass. + * + * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. + **/ + void set_charge_spin(const std::vector& charge_spin) override; void get_type_map(std::string& type_map); bool is_aparam_nall() const { assert(inited); @@ -414,10 +442,19 @@ class DeepPotPTExpt : public DeepPotBackend { const int nall_nodes, const std::int64_t edge_storage); bool inited; - int ntypes; - int dfparam; - int daparam; - int dchgspin; + // Every width below is a property of the loaded archive. They read as zero + // until ``init`` has run, so that a query on an uninitialised backend + // answers "none" rather than whatever the allocation happened to hold. + int ntypes = 0; + int dfparam = 0; + int daparam = 0; + // Conditioning width of the compiled forward's argument list. Zero for a + // model that carries no condition, and also for a compressed one, whose + // condition lives in frozen tables rather than in an input. Every gate that + // decides whether to hand the forward a condition tensor reads this. + int dchgspin = 0; + // Width of a charge state this model can be given; see ``dim_chg_spin()``. + int settable_chgspin = 0; bool aparam_nall; bool has_default_fparam_; std::vector default_fparam_; @@ -473,6 +510,11 @@ class DeepPotPTExpt : public DeepPotBackend { torch::Tensor pair_exclude_table_; std::unique_ptr with_comm_tempfile_; std::unique_ptr with_comm_loader; + // The charge/spin condition a compressed descriptor folded into the + // constants of its lower, re-runnable so that a condition chosen at runtime + // can be served. Null for an uncompressed model, which reads its condition + // as an ordinary input and needs no rebuild. + std::unique_ptr charge_state_fold_; /** * @brief Multi-frame loop for standalone compute (no nlist). diff --git a/source/api_cc/include/DeepSpin.h b/source/api_cc/include/DeepSpin.h index 14d3d685b8..bf57849b14 100644 --- a/source/api_cc/include/DeepSpin.h +++ b/source/api_cc/include/DeepSpin.h @@ -168,6 +168,27 @@ class DeepSpinBackend : public DeepBaseModelBackend { **/ virtual int dim_chg_spin() const { return 0; } + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * + * An override is needed only where the condition has to be folded into state + * that is built ahead of the evaluations using it, as in a compressed model + * whose tables are specialized to one condition. A backend that reads the + * condition as an ordinary per-call input has nothing to install, so the + * default is a no-op rather than an error: the condition still reaches such + * a backend on every evaluation, through the charge_spin argument of + * computew(). The request is refused only by a model that carries no + * charge/spin conditioning at all, which no route can honour. + * + * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. + **/ + virtual void set_charge_spin(const std::vector& charge_spin) { + if (dim_chg_spin() == 0) { + throw deepmd::deepmd_exception( + "this model does not support a charge/spin condition"); + } + } + // charge_spin-aware computew overloads. Default implementations call the // existing pure-virtual overloads (ignoring charge_spin) so that backends // that do not support charge/spin do not need any changes. DeepSpinPTExpt @@ -575,6 +596,12 @@ class DeepSpin : public DeepBaseModel { **/ int dim_chg_spin() const; + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. + **/ + void set_charge_spin(const std::vector& charge_spin); + /** * @brief Get the per-type use_spin flags. * @return A vector of booleans indicating which atom types have spin enabled. @@ -679,6 +706,18 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { return numb_models > 0 ? dps[0]->dim_chg_spin() : 0; }; + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * Applied to every model, so that the deviation is taken between models + * under the same condition. + * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. + **/ + void set_charge_spin(const std::vector& charge_spin) { + for (unsigned ii = 0; ii < dps.size(); ++ii) { + dps[ii]->set_charge_spin(charge_spin); + } + }; + /** * @brief Evaluate the energy, force and virial by using these DP spin models. * @param[out] all_ener The system energies of all models. diff --git a/source/api_cc/include/DeepSpinPTExpt.h b/source/api_cc/include/DeepSpinPTExpt.h index 500c492607..778d86ce39 100644 --- a/source/api_cc/include/DeepSpinPTExpt.h +++ b/source/api_cc/include/DeepSpinPTExpt.h @@ -105,6 +105,17 @@ class DeepSpinPTExpt : public DeepSpinBackend { assert(inited); return dchgspin; }; + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * + * It becomes the condition of every later forward pass that is not given + * one explicitly. This backend serves models whose lower reads the + * condition as an ordinary input, so that stored value is the whole + * mechanism. + * + * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. + **/ + void set_charge_spin(const std::vector& charge_spin) override; void get_type_map(std::string& type_map); bool is_aparam_nall() const { assert(inited); diff --git a/source/api_cc/include/NativeSpinPTExpt.h b/source/api_cc/include/NativeSpinPTExpt.h index 6f990b24ac..4830bcd67f 100644 --- a/source/api_cc/include/NativeSpinPTExpt.h +++ b/source/api_cc/include/NativeSpinPTExpt.h @@ -16,6 +16,12 @@ #include "DeepSpin.h" +// Forward-declare to keep the private header out of the public one. Defined in +// commonPTExpt.h. +namespace deepmd::ptexpt { +class ChargeStateFold; +} + namespace torch::inductor { class AOTIModelPackageLoader; } @@ -41,8 +47,9 @@ struct CanonicalGraphTensorPack; * * - ``lower_input_kind == "graph"``, the general NeighborGraph ABI * (``forward_lower_graph_exportable``): the ten topology tensors, the - * per-node moment at positional index 10, then the conditional frame and - * atomic parameter tail. Any graph-lower descriptor can be frozen this way. + * per-node moment at positional index 10, then the conditional tail of + * frame parameter, atomic parameter and charge/spin condition. Any + * graph-lower descriptor can be frozen this way. * - ``lower_input_kind == "dpa4c_canonical"``, the compact deployment ABI * (``forward_lower_canonical_graph_exportable``): the eight dual-CSR graph * tensors -- uint32 topology and float32 edge vectors -- and the moment at @@ -100,10 +107,37 @@ class NativeSpinPTExpt : public DeepSpinBackend { assert(inited); return daparam; }; - // Charge/spin conditioning is rejected when a native-spin model is built, - // so no archive this backend accepts carries a non-zero width; ``init`` - // enforces it. - int dim_chg_spin() const override { return 0; }; + /** + * @brief The width of a charge/spin condition this model accepts. + * + * This is the width a caller names a condition with, which is not in + * general the width of the conditioning input of the compiled forward: + * compression folds the condition into frozen tables and so removes it from + * the argument list, leaving a model that still serves a condition through + * the fold shipped beside the inference lower. + **/ + int dim_chg_spin() const override { + assert(inited); + return settable_chgspin; + }; + /** + * @brief Fix the charge/spin condition served for the rest of the run. + * + * The condition reaches the model by one of two routes, and this sets both + * so that the caller does not depend on how the model was frozen. It + * becomes the condition of every later forward pass that is not given one + * explicitly. A compressed descriptor additionally carries the condition + * inside frozen tables that the compiled lower holds as constants; when the + * archive ships the fold that rebuilds them, it runs here and the resulting + * tables are written over those constants. + * + * Intended to be called once, before inference, since overwriting the + * constants of a loaded module is not safe to interleave with a forward + * pass. + * + * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. + **/ + void set_charge_spin(const std::vector& charge_spin) override; void get_type_map(std::string& type_map); bool is_aparam_nall() const { return false; }; bool has_default_fparam() const { @@ -174,6 +208,74 @@ class NativeSpinPTExpt : public DeepSpinBackend { const std::vector& aparam, const bool atomic); + // Charge/spin-aware overloads. This backend serves the condition in force + // rather than marshalling one per call, so a condition named here is + // checked against that state and rejected when it names another; the + // inherited defaults would drop it without a word. An empty condition + // selects the state in force. + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + /** * @brief Fully device-resident inference on a compact canonical graph. * @@ -279,13 +381,15 @@ class NativeSpinPTExpt : public DeepSpinBackend { * @brief Run the NeighborGraph native-spin forward. * * Positional order: the ten NeighborGraph tensors, the per-node moment at - * index 10, then the conditional tail -- the frame parameter and the atomic - * parameter, each present only when the model declares a non-zero width. + * index 10, then the conditional tail -- the frame parameter, the atomic + * parameter and the charge/spin condition, each present only when the model + * declares a non-zero width. */ std::vector run_model_graph(const GraphTensorPack& graph, const torch::Tensor& spin, const torch::Tensor& fparam, - const torch::Tensor& aparam); + const torch::Tensor& aparam, + const torch::Tensor& charge_spin); /** * @brief Apply model-level pair exclusion, canonicalize the payload and run. @@ -297,6 +401,10 @@ class NativeSpinPTExpt : public DeepSpinBackend { * forward. The returned map holds the artifact's public output keys with the * per-atom virial in its ``(N, 9)`` layout. * + * The charge/spin condition is not marshalled per call: the NeighborGraph + * forward reads the condition currently in force, and the compact one + * carries it in its frozen tables. + * * @param[in,out] graph Graph payload for ``node_count`` nodes; consumed in * place by the canonicalization. * @param[in] node_count Number of graph nodes. @@ -327,12 +435,26 @@ class NativeSpinPTExpt : public DeepSpinBackend { void translate_error(std::function f); bool inited; - int ntypes; - int ntypes_spin; - int dfparam; - int daparam; + // Every width below is a property of the loaded archive. They read as zero + // until ``init`` has run, so that a query on an uninitialised backend + // answers "none" rather than whatever the allocation happened to hold. + int ntypes = 0; + int ntypes_spin = 0; + int dfparam = 0; + int daparam = 0; + // Conditioning width of the compiled forward's argument list. Zero for a + // model that carries no condition, and also for a compressed one, whose + // condition lives in frozen tables rather than in an input. Every gate that + // decides whether to hand the forward a condition tensor reads this. + int dchgspin = 0; + // Width of a charge state this model can be given; see ``dim_chg_spin()``. + int settable_chgspin = 0; bool has_default_fparam_; std::vector default_fparam_; + // The condition served by every forward pass that is not given one + // explicitly, initialised from the archive and replaced by + // ``set_charge_spin``. + std::vector default_chg_spin_; double rcut; int gpu_id; bool gpu_enabled; @@ -362,6 +484,11 @@ class NativeSpinPTExpt : public DeepSpinBackend { at::Tensor edge_index_tensor; // node-space edges (folded or extended) at::Tensor edge_index_ext_tensor; // extended-atom edges, for the geometry std::unique_ptr loader; + // The charge/spin condition a compressed descriptor folded into the + // constants of its lower, re-runnable so that a condition chosen at runtime + // can be served. Null for an uncompressed model, which reads its condition + // as an ordinary input and needs no rebuild. + std::unique_ptr charge_state_fold_; }; } // namespace deepmd diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 1e45a24438..7331eab8b2 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -196,6 +196,18 @@ inline EdgeTensorPack createEdgeTensors( const bool with_geometry = true, const std::vector* row_centers = nullptr, const bool fold_to_local = true) { + // Folding reads an owner for every extended atom, so a mapping shorter than + // that is a caller error rather than a degraded input: indexing it would be + // out of bounds, and the out-of-range owners it appears to yield would drop + // the corresponding edges one by one, leaving a quietly incomplete graph. + if (fold_to_local && mapping.size() < static_cast(nall)) { + throw deepmd::deepmd_exception( + "folding ghost neighbours onto their local owners needs an owner for " + "each of the " + + std::to_string(nall) + " extended atoms, but the mapping holds " + + std::to_string(mapping.size()) + + "; under LAMMPS this is what 'atom_modify map yes' supplies"); + } std::vector src; std::vector dst; std::vector src_ext; @@ -235,8 +247,19 @@ inline EdgeTensorPack createEdgeTensors( std::int64_t src_node; if (fold_to_local) { const std::int64_t src_local = mapping[static_cast(jj)]; + // Folding is single-domain, where every extended atom has an owner + // among the local ones. An owner outside that range therefore marks a + // mapping that was never filled, not a neighbour to skip: skipping + // would discard every ghost edge and leave a graph that is quietly + // missing the whole halo. if (src_local < 0 || src_local >= nloc) { - continue; + throw deepmd::deepmd_exception( + "extended atom " + std::to_string(jj) + " of " + + std::to_string(nall) + " maps to owner " + + std::to_string(src_local) + ", which is not one of the " + + std::to_string(nloc) + + " local atoms; under LAMMPS an owner for every extended atom is " + "what 'atom_modify map yes' supplies"); } src_node = src_local; } else { diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index dd9b57e2e3..3373d6fe71 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -778,6 +778,10 @@ bool DeepPot::uses_canonical_graph_inference() const { int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); } +void DeepPot::set_charge_spin(const std::vector& charge_spin) { + dp->set_charge_spin(charge_spin); +} + DeepPotModelDevi::DeepPotModelDevi() { inited = false; numb_models = 0; diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 979c36cbfe..b1792066c8 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -14,6 +14,7 @@ #include #include #include +#include #include "SimulationRegion.h" #include "common.h" @@ -23,6 +24,7 @@ #include "errors.h" #include "neighbor_list.h" +using deepmd::ptexpt::check_call_charge_spin; using deepmd::ptexpt::parse_json; using deepmd::ptexpt::read_default_chg_spin; using deepmd::ptexpt::read_zip_entry; @@ -190,6 +192,10 @@ void DeepPotPTExpt::init(const std::string& model, dchgspin = metadata.obj_val.count("dim_chg_spin") ? metadata["dim_chg_spin"].as_int() : 0; + // A model whose lower reads the condition as an input accepts exactly that + // condition; the charge-state fold, loaded below, widens this for a + // compressed model, whose lower has no conditioning input at all. + settable_chgspin = dchgspin; aparam_nall = false; // pt_expt models use nloc for aparam if (metadata.obj_val.count("has_default_fparam")) { has_default_fparam_ = metadata["has_default_fparam"].as_bool(); @@ -364,6 +370,31 @@ void DeepPotPTExpt::init(const std::string& model, } } + // Charge-state fold. Unlike the with-comm artifact, a failure here is not + // a degraded mode to defer: the metadata field is the archive's claim that + // the fold ships with it, so an archive that declares the constants but + // cannot supply the fold is malformed. + charge_state_fold_ = deepmd::ptexpt::ChargeStateFold::load( + model, metadata, gpu_enabled, gpu_id); + if (charge_state_fold_) { + // The condition of a compressed model reaches the compiled lower only + // through the constants the fold rebuilds, so the argument list carries + // none and ``dchgspin`` is zero. What the model accepts is the state the + // snapshot was frozen against, which is also the layout the fold consumes. + settable_chgspin = + metadata.obj_val.count("default_chg_spin") + ? static_cast(metadata["default_chg_spin"].as_array().size()) + : 0; + if (settable_chgspin == 0) { + throw deepmd::deepmd_exception( + "the archive ships a charge-state fold but names no " + "default_chg_spin, so the width of a charge state is unknown"); + } + // The constants were frozen against the archive's own charge state, so + // that is the state in force until ``set_charge_spin`` installs another. + default_chg_spin_ = read_default_chg_spin(metadata, settable_chgspin); + } + int num_intra_nthreads, num_inter_nthreads; get_env_nthreads(num_intra_nthreads, num_inter_nthreads); if (num_inter_nthreads) { @@ -384,6 +415,33 @@ void DeepPotPTExpt::init(const std::string& model, DeepPotPTExpt::~DeepPotPTExpt() {} +void DeepPotPTExpt::set_charge_spin(const std::vector& charge_spin) { + assert(inited); + if (settable_chgspin == 0) { + throw deepmd::deepmd_exception( + "this model was not frozen with a charge/spin condition"); + } + if (static_cast(charge_spin.size()) != settable_chgspin) { + throw deepmd::deepmd_exception("the charge/spin condition carries " + + std::to_string(charge_spin.size()) + + " values but the model expects " + + std::to_string(settable_chgspin)); + } + // Route one: the condition of every later forward pass that is not given + // one explicitly. This is the whole mechanism for an uncompressed model, + // which reads the condition as an ordinary input. + default_chg_spin_ = charge_spin; + // Route two: a compressed descriptor has folded the condition into frozen + // tables that the lower holds as constants, so serving another condition + // means rebuilding those tables and writing them over the constants. + if (charge_state_fold_) { + charge_state_fold_->apply(charge_spin, + gpu_enabled ? torch::Device(torch::kCUDA, gpu_id) + : torch::Device(torch::kCPU), + *loader); + } +} + std::vector DeepPotPTExpt::run_model( const torch::Tensor& coord, const torch::Tensor& atype, @@ -667,6 +725,10 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, "and is off by default for .pt2. To enable it, regenerate with: " "dp convert-backend --atomic-virial INPUT.pth OUTPUT.pt2"); } + // A single-frame call names at most one charge state, and only one this + // model can serve. + check_call_charge_spin(charge_spin, /*nframes=*/1, settable_chgspin, + dchgspin > 0, default_chg_spin_); torch::Device device(torch::kCUDA, gpu_id); if (!gpu_enabled) { device = torch::Device(torch::kCPU); @@ -725,6 +787,18 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, bool multi_rank = (lmp_list.nprocs > 1); bool atom_map_present = (lmp_list.mapping != nullptr); bool use_with_comm = has_comm_artifact_ && multi_rank; + // Whether the edge topology folds ghost neighbours onto their local owners, + // which reads an owner for every extended atom out of ``mapping_``. The + // edge lower folds unless the with-comm artifact carries ghost features + // across ranks; the graph and canonical lowers fold on a single rank and + // keep the extended node set under domain decomposition, where ghost forces + // reverse-communicate to their owners instead. The dense lower builds no + // edge topology and therefore never folds. + const bool fold_to_local = + lower_input_is_edge_ + ? !use_with_comm + : ((lower_input_is_graph_ || lower_input_is_canonical_) && + !multi_rank); // NeighborGraph multi-rank dispatch: // - NON-message-passing (dpa1, se_e2_a, ...): the SAME single-rank graph // .pt2 runs on the EXTENDED region (fold_to_local=false; ghosts are @@ -766,6 +840,21 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, "callers must set inlist.mapping explicitly before compute()."); } } + // Folding resolves every real extended atom to a local owner, so it needs + // the atom map for the same reason the ghost-feature gather above does, and + // for every lower that builds an edge topology rather than only for a + // message-passing one. Without the map the owner table below degrades to + // the identity, whose ghost entries are not local atoms; establishing the + // precondition here reports it once, ahead of any tensor, instead of + // leaving it to the per-edge lookup inside the topology build. + if (fold_to_local && nghost_real > 0 && !atom_map_present) { + throw deepmd::deepmd_exception( + "This .pt2 lower folds ghost neighbours onto their local owners, " + "which needs an owner for each of the " + + std::to_string(nghost_real) + + " ghost atoms: add `atom_modify map yes` to the LAMMPS input, or, as " + "a C++ API caller, set inlist.mapping before compute()."); + } // LAMMPS sets ago=0 on every nlist rebuild (neighbor rebuild, re-partition, // atom exchange between subdomains), so `ago > 0` implies the cached @@ -791,11 +880,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, .clone() .to(device); } else { - // Identity fallback. The fail-fast above guarantees we only + // Identity fallback. The fail-fasts above guarantee we only // reach this branch when one of these is true: - // - The model is non-message-passing (mapping is unused). - // - ``nghost == 0`` (no ghosts to gather, identity is trivially - // correct). + // - No real ghost exists, so the identity is the owner table. + // - The topology keeps the extended node set (``!fold_to_local``) + // and the model is non-message-passing, leaving this tensor + // unread. // - ``use_with_comm`` is true (the with-comm graph fills ghost // features via border_op and ignores this tensor for ghost // gather — see deepmd/pt_expt/descriptor/ @@ -823,7 +913,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, const auto edge_tensors = createEdgeTensors( nlist_data.jlist, dcoord, mapping_, nloc, nall_real, device, /*with_geometry=*/false, /*row_centers=*/&nlist_data.ilist, - /*fold_to_local=*/!use_with_comm); + fold_to_local); edge_index_tensor = edge_tensors.edge_index; edge_index_ext_tensor = edge_tensors.edge_index_ext; } else if (lower_input_is_graph_ || lower_input_is_canonical_) { @@ -833,7 +923,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, const auto edge_tensors = createEdgeTensors( nlist_data.jlist, dcoord, mapping_, nloc, nall_real, device, /*with_geometry=*/false, /*row_centers=*/&nlist_data.ilist, - /*fold_to_local=*/!multi_rank); + fold_to_local); edge_index_tensor = edge_tensors.edge_index; edge_index_ext_tensor = edge_tensors.edge_index_ext; } else { @@ -873,37 +963,23 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, at::Tensor aparam_tensor = make_aparam_tensor(aparam_, nloc, daparam, device); - // Build charge_spin tensor: use runtime value when provided, fall back to - // default_chg_spin_ stored in the .pt2 metadata. + // Build charge_spin tensor: the condition supplied with the call when there + // is one, otherwise the state in force. at::Tensor charge_spin_tensor; if (dchgspin > 0) { - auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); - if (!charge_spin.empty()) { - // Single-frame path: charge_spin must hold exactly dim_chg_spin values. - if (static_cast(charge_spin.size()) != dchgspin) { - throw deepmd::deepmd_exception( - "charge_spin has " + std::to_string(charge_spin.size()) + - " values but the model expects dim_chg_spin=" + - std::to_string(dchgspin) + "."); - } - charge_spin_tensor = - torch::from_blob(const_cast(charge_spin.data()), - {1, static_cast(charge_spin.size())}, - dbl_options) - .clone() - .to(device); - } else if (!default_chg_spin_.empty()) { - charge_spin_tensor = - torch::from_blob(const_cast(default_chg_spin_.data()), - {1, dchgspin}, dbl_options) - .clone() - .to(device); - } else { + const std::vector& condition = + charge_spin.empty() ? default_chg_spin_ : charge_spin; + if (condition.empty()) { throw deepmd::deepmd_exception( "charge_spin is empty and no default_chg_spin is available in the " ".pt2 metadata. Provide charge_spin explicitly or regenerate the " "model with a default charge/spin value."); } + charge_spin_tensor = + torch::from_blob(const_cast(condition.data()), {1, dchgspin}, + torch::TensorOptions().dtype(torch::kFloat64)) + .clone() + .to(device); } // ``use_with_comm`` was computed earlier alongside the fail-fast @@ -1388,6 +1464,10 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, coord, atype, box, fparam, aparam, charge_spin, atomic); return; } + // A single-frame call names at most one charge state, and only one this + // model can serve. + check_call_charge_spin(charge_spin, /*nframes=*/1, settable_chgspin, + dchgspin > 0, default_chg_spin_); // The .pt2 model only contains forward_common_lower, which requires // nlist as input. We must build the nlist in C++ and fold back the // extended-region outputs to local atoms. @@ -1528,37 +1608,23 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, at::Tensor aparam_tensor = make_aparam_tensor(aparam, natoms, daparam, device); - // Build charge_spin tensor: use runtime value when provided, fall back to - // default_chg_spin_ stored in the .pt2 metadata. + // Build charge_spin tensor: the condition supplied with the call when there + // is one, otherwise the state in force. at::Tensor charge_spin_tensor; if (dchgspin > 0) { - auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); - if (!charge_spin.empty()) { - // Single-frame path: charge_spin must hold exactly dim_chg_spin values. - if (static_cast(charge_spin.size()) != dchgspin) { - throw deepmd::deepmd_exception( - "charge_spin has " + std::to_string(charge_spin.size()) + - " values but the model expects dim_chg_spin=" + - std::to_string(dchgspin) + "."); - } - charge_spin_tensor = - torch::from_blob(const_cast(charge_spin.data()), - {1, static_cast(charge_spin.size())}, - dbl_options) - .clone() - .to(device); - } else if (!default_chg_spin_.empty()) { - charge_spin_tensor = - torch::from_blob(const_cast(default_chg_spin_.data()), - {1, dchgspin}, dbl_options) - .clone() - .to(device); - } else { + const std::vector& condition = + charge_spin.empty() ? default_chg_spin_ : charge_spin; + if (condition.empty()) { throw deepmd::deepmd_exception( "charge_spin is empty and no default_chg_spin is available in the " ".pt2 metadata. Provide charge_spin explicitly or regenerate the " "model with a default charge/spin value."); } + charge_spin_tensor = + torch::from_blob(const_cast(condition.data()), {1, dchgspin}, + torch::TensorOptions().dtype(torch::kFloat64)) + .clone() + .to(device); } // 5. Run the .pt2 model @@ -1696,21 +1762,11 @@ void DeepPotPTExpt::compute_nframes(ENERGYVTYPE& ener, static_cast(natoms) * daparam, false); const FrameParameterLayout fparam_layout = resolve_frame_parameter_layout( "fparam", fparam.size(), nframes, dfparam, true); - // charge_spin may be empty (default fallback), a single dim_chg_spin vector - // (broadcast to all frames), or nframes * dim_chg_spin (per-frame). Reject - // anything else up-front to avoid out-of-range slicing in the loop. - if (!charge_spin.empty()) { - size_t s_dcsp = static_cast(dchgspin); - if (charge_spin.size() != s_dcsp && - charge_spin.size() != s_dcsp * static_cast(nframes)) { - throw deepmd::deepmd_exception( - "charge_spin has " + std::to_string(charge_spin.size()) + - " values but the model expects dim_chg_spin=" + - std::to_string(dchgspin) + " (per frame) or " + - std::to_string(dchgspin * nframes) + " (for " + - std::to_string(nframes) + " frames)."); - } - } + // charge_spin may be empty (the state in force), one charge state + // (broadcast to every frame), or one per frame. Reject anything else up + // front to avoid out-of-range slicing in the loop. + check_call_charge_spin(charge_spin, nframes, settable_chgspin, dchgspin > 0, + default_chg_spin_); ener.clear(); force.clear(); virial.clear(); @@ -1744,9 +1800,9 @@ void DeepPotPTExpt::compute_nframes(ENERGYVTYPE& ener, } std::vector frame_chg_spin; if (!charge_spin.empty()) { - size_t s_dcsp = static_cast(dchgspin); + size_t s_dcsp = static_cast(settable_chgspin); if (charge_spin.size() == s_dcsp) { - // single charge/spin vector broadcast to every frame + // one charge state broadcast to every frame frame_chg_spin = charge_spin; } else { frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, @@ -1910,21 +1966,11 @@ void DeepPotPTExpt::compute_mixed_type_impl( static_cast(natoms) * daparam, false); const FrameParameterLayout fparam_layout = resolve_frame_parameter_layout( "fparam", fparam.size(), nframes, dfparam, true); - // charge_spin may be empty (default fallback), a single dim_chg_spin vector - // (broadcast to all frames), or nframes * dim_chg_spin (per-frame). Reject - // anything else up-front to avoid out-of-range slicing in the loop. - if (!charge_spin.empty()) { - size_t s_dcsp = static_cast(dchgspin); - if (charge_spin.size() != s_dcsp && - charge_spin.size() != s_dcsp * static_cast(nframes)) { - throw deepmd::deepmd_exception( - "charge_spin has " + std::to_string(charge_spin.size()) + - " values but the model expects dim_chg_spin=" + - std::to_string(dchgspin) + " (per frame) or " + - std::to_string(dchgspin * nframes) + " (for " + - std::to_string(nframes) + " frames)."); - } - } + // charge_spin may be empty (the state in force), one charge state + // (broadcast to every frame), or one per frame. Reject anything else up + // front to avoid out-of-range slicing in the loop. + check_call_charge_spin(charge_spin, nframes, settable_chgspin, dchgspin > 0, + default_chg_spin_); ener.clear(); force.clear(); virial.clear(); @@ -1960,9 +2006,9 @@ void DeepPotPTExpt::compute_mixed_type_impl( } std::vector frame_chg_spin; if (!charge_spin.empty()) { - size_t s_dcsp = static_cast(dchgspin); + size_t s_dcsp = static_cast(settable_chgspin); if (charge_spin.size() == s_dcsp) { - // single charge/spin vector broadcast to every frame + // one charge state broadcast to every frame frame_chg_spin = charge_spin; } else { frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, diff --git a/source/api_cc/src/DeepSpin.cc b/source/api_cc/src/DeepSpin.cc index 28731ab0d8..e52d6a657b 100644 --- a/source/api_cc/src/DeepSpin.cc +++ b/source/api_cc/src/DeepSpin.cc @@ -458,6 +458,10 @@ template void DeepSpin::compute(std::vector& dener, int DeepSpin::dim_chg_spin() const { return dp->dim_chg_spin(); } +void DeepSpin::set_charge_spin(const std::vector& charge_spin) { + dp->set_charge_spin(charge_spin); +} + std::vector DeepSpin::get_use_spin() const { if (dp) { return dp->get_use_spin(); diff --git a/source/api_cc/src/DeepSpinPTExpt.cc b/source/api_cc/src/DeepSpinPTExpt.cc index 958b240ead..fbbf46c689 100644 --- a/source/api_cc/src/DeepSpinPTExpt.cc +++ b/source/api_cc/src/DeepSpinPTExpt.cc @@ -20,6 +20,7 @@ #include "errors.h" #include "neighbor_list.h" +using deepmd::ptexpt::check_call_charge_spin; using deepmd::ptexpt::parse_json; using deepmd::ptexpt::read_default_chg_spin; using deepmd::ptexpt::read_zip_entry; @@ -290,6 +291,21 @@ void DeepSpinPTExpt::init(const std::string& model, DeepSpinPTExpt::~DeepSpinPTExpt() {} +void DeepSpinPTExpt::set_charge_spin(const std::vector& charge_spin) { + assert(inited); + if (dchgspin == 0) { + throw deepmd::deepmd_exception( + "this model was not frozen with a charge/spin condition"); + } + if (static_cast(charge_spin.size()) != dchgspin) { + throw deepmd::deepmd_exception("the charge/spin condition carries " + + std::to_string(charge_spin.size()) + + " values but the model expects " + + std::to_string(dchgspin)); + } + default_chg_spin_ = charge_spin; +} + std::vector DeepSpinPTExpt::run_model( const torch::Tensor& coord, const torch::Tensor& atype, @@ -595,6 +611,10 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, "and is off by default for .pt2. To enable it, regenerate with: " "dp convert-backend --atomic-virial INPUT.pth OUTPUT.pt2"); } + // A single-frame call names at most one charge state, and only one this + // model can serve. + check_call_charge_spin(charge_spin, /*nframes=*/1, dchgspin, dchgspin > 0, + default_chg_spin_); torch::Device device(torch::kCUDA, gpu_id); if (!gpu_enabled) { device = torch::Device(torch::kCPU); @@ -740,6 +760,21 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, } } + // Folding ghost neighbours onto their local owners reads an owner for every + // extended atom out of the mapping, which the preceding matrix guarantees + // only for a message-passing model. Without a mapping the owner table falls + // back to the identity, whose ghost entries are not local atoms; + // establishing the precondition here reports it once, ahead of any tensor, + // instead of leaving it to the per-edge lookup inside the topology build. + if (!use_with_comm && nghost_real > 0 && !atom_map_present) { + throw deepmd::deepmd_exception( + "This .pt2 lower folds ghost neighbours onto their local owners, " + "which needs an owner for each of the " + + std::to_string(nghost_real) + + " ghost atoms: add `atom_modify map yes` to the LAMMPS input, or, as " + "a C++ API caller, set inlist.mapping before compute()."); + } + // LAMMPS sets ago=0 on every nlist rebuild, so ago>0 implies the cached // mapping and nlist tensors are still valid — see DeepPotPTExpt.cc for // the same rationale. @@ -868,38 +903,23 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, aparam_tensor = torch::zeros({0}, options).to(device); } - // Build charge_spin tensor: use the runtime value when provided, fall back - // to default_chg_spin_ stored in the .pt2 metadata. Mirrors - // DeepPotPTExpt::compute -- these spin paths are single-frame, so the - // runtime vector must hold exactly dim_chg_spin values. + // Build charge_spin tensor: the condition supplied with the call when there + // is one, otherwise the state in force. at::Tensor charge_spin_tensor; if (dchgspin > 0) { - auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); - if (!charge_spin.empty()) { - if (static_cast(charge_spin.size()) != dchgspin) { - throw deepmd::deepmd_exception( - "charge_spin has " + std::to_string(charge_spin.size()) + - " values but the model expects dim_chg_spin=" + - std::to_string(dchgspin) + "."); - } - charge_spin_tensor = - torch::from_blob(const_cast(charge_spin.data()), - {1, static_cast(charge_spin.size())}, - dbl_options) - .clone() - .to(device); - } else if (!default_chg_spin_.empty()) { - charge_spin_tensor = - torch::from_blob(const_cast(default_chg_spin_.data()), - {1, dchgspin}, dbl_options) - .clone() - .to(device); - } else { + const std::vector& condition = + charge_spin.empty() ? default_chg_spin_ : charge_spin; + if (condition.empty()) { throw deepmd::deepmd_exception( "charge_spin is empty and no default_chg_spin is available in the " ".pt2 metadata. Provide charge_spin explicitly or regenerate the " "model with a default charge/spin value."); } + charge_spin_tensor = + torch::from_blob(const_cast(condition.data()), {1, dchgspin}, + torch::TensorOptions().dtype(torch::kFloat64)) + .clone() + .to(device); } // Phase 4 dispatch: route to with-comm artifact in multi-rank mode. @@ -1371,6 +1391,10 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, "and is off by default for .pt2. To enable it, regenerate with: " "dp convert-backend --atomic-virial INPUT.pth OUTPUT.pt2"); } + // A single-frame call names at most one charge state, and only one this + // model can serve. + check_call_charge_spin(charge_spin, /*nframes=*/1, dchgspin, dchgspin > 0, + default_chg_spin_); int natoms = atype.size(); torch::Device device(torch::kCUDA, gpu_id); @@ -1543,38 +1567,23 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, aparam_tensor = torch::zeros({0}, options).to(device); } - // Build charge_spin tensor: use the runtime value when provided, fall back - // to default_chg_spin_ stored in the .pt2 metadata. Mirrors - // DeepPotPTExpt::compute -- these spin paths are single-frame, so the - // runtime vector must hold exactly dim_chg_spin values. + // Build charge_spin tensor: the condition supplied with the call when there + // is one, otherwise the state in force. at::Tensor charge_spin_tensor; if (dchgspin > 0) { - auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); - if (!charge_spin.empty()) { - if (static_cast(charge_spin.size()) != dchgspin) { - throw deepmd::deepmd_exception( - "charge_spin has " + std::to_string(charge_spin.size()) + - " values but the model expects dim_chg_spin=" + - std::to_string(dchgspin) + "."); - } - charge_spin_tensor = - torch::from_blob(const_cast(charge_spin.data()), - {1, static_cast(charge_spin.size())}, - dbl_options) - .clone() - .to(device); - } else if (!default_chg_spin_.empty()) { - charge_spin_tensor = - torch::from_blob(const_cast(default_chg_spin_.data()), - {1, dchgspin}, dbl_options) - .clone() - .to(device); - } else { + const std::vector& condition = + charge_spin.empty() ? default_chg_spin_ : charge_spin; + if (condition.empty()) { throw deepmd::deepmd_exception( "charge_spin is empty and no default_chg_spin is available in the " ".pt2 metadata. Provide charge_spin explicitly or regenerate the " "model with a default charge/spin value."); } + charge_spin_tensor = + torch::from_blob(const_cast(condition.data()), {1, dchgspin}, + torch::TensorOptions().dtype(torch::kFloat64)) + .clone() + .to(device); } // 5. Run the .pt2 model: native spin uses the energy edge ABI plus the diff --git a/source/api_cc/src/NativeSpinPTExpt.cc b/source/api_cc/src/NativeSpinPTExpt.cc index 18adbf64c1..bfa46bc71e 100644 --- a/source/api_cc/src/NativeSpinPTExpt.cc +++ b/source/api_cc/src/NativeSpinPTExpt.cc @@ -21,7 +21,9 @@ #include "errors.h" #include "neighbor_list.h" +using deepmd::ptexpt::check_call_charge_spin; using deepmd::ptexpt::parse_json; +using deepmd::ptexpt::read_default_chg_spin; using deepmd::ptexpt::read_zip_entry; using namespace deepmd; @@ -119,6 +121,35 @@ torch::Tensor make_aparam_tensor(const std::vector& aparam, return deepmd::extend_graph_aparam(owned, node_count, nloc, dim_aparam); } +/** + * @brief Build the charge/spin input of the conditional graph tail. + * + * The artifact consumes the condition in double precision, shaped + * ``(1, dim_chg_spin)``. A zero width means the forward has no such slot -- + * either because the model carries no condition at all, or because + * compression moved it into frozen tables -- so the returned tensor is + * undefined and never marshalled. + */ +torch::Tensor make_chg_spin_tensor(const std::vector& charge_spin, + const int dim_chg_spin, + const torch::Device& device) { + if (dim_chg_spin == 0) { + return torch::Tensor(); + } + if (static_cast(charge_spin.size()) != dim_chg_spin) { + throw deepmd::deepmd_exception( + "the charge/spin condition holds " + + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dim_chg_spin) + "."); + } + return torch::from_blob(const_cast(charge_spin.data()), + {1, dim_chg_spin}, + torch::TensorOptions().dtype(torch::kFloat64)) + .clone() + .to(device); +} + } // namespace void NativeSpinPTExpt::translate_error(std::function f) { @@ -211,16 +242,14 @@ void NativeSpinPTExpt::init(const std::string& model, const int declared_chg_spin = metadata.obj_val.count("dim_chg_spin") ? metadata["dim_chg_spin"].as_int() : 0; - // Charge/spin FiLM conditioning is rejected when a native-spin model is - // built, so neither schema reserves a slot for it. - if (declared_chg_spin > 0) { - throw deepmd::deepmd_exception( - "native spin does not combine with charge/spin conditioning, but this " - "archive declares dim_chg_spin=" + - std::to_string(declared_chg_spin) + "."); - } dfparam = declared_fparam; daparam = declared_aparam; + dchgspin = declared_chg_spin; + // A model whose lower reads the condition as an input accepts exactly that + // condition; the charge-state fold, loaded below, widens this for a + // compressed model, whose lower has no conditioning input at all. + settable_chgspin = dchgspin; + default_chg_spin_ = read_default_chg_spin(metadata, dchgspin); has_default_fparam_ = metadata.obj_val.count("has_default_fparam") && metadata["has_default_fparam"].as_bool(); default_fparam_.clear(); @@ -321,6 +350,31 @@ void NativeSpinPTExpt::init(const std::string& model, } } + // Charge-state fold. Unlike the with-comm artifact, a failure here is not + // a degraded mode to defer: the metadata field is the archive's claim that + // the fold ships with it, so an archive that declares the constants but + // cannot supply the fold is malformed. + charge_state_fold_ = deepmd::ptexpt::ChargeStateFold::load( + model, metadata, gpu_enabled, gpu_id); + if (charge_state_fold_) { + // The condition of a compressed model reaches the compiled lower only + // through the constants the fold rebuilds, so the argument list carries + // none and ``dchgspin`` is zero. What the model accepts is the state the + // snapshot was frozen against, which is also the layout the fold consumes. + settable_chgspin = + metadata.obj_val.count("default_chg_spin") + ? static_cast(metadata["default_chg_spin"].as_array().size()) + : 0; + if (settable_chgspin == 0) { + throw deepmd::deepmd_exception( + "the archive ships a charge-state fold but names no " + "default_chg_spin, so the width of a charge state is unknown"); + } + // The constants were frozen against the archive's own charge state, so + // that is the state in force until ``set_charge_spin`` installs another. + default_chg_spin_ = read_default_chg_spin(metadata, settable_chgspin); + } + int num_intra_nthreads, num_inter_nthreads; get_env_nthreads(num_intra_nthreads, num_inter_nthreads); if (num_inter_nthreads) { @@ -339,6 +393,33 @@ void NativeSpinPTExpt::init(const std::string& model, inited = true; } +void NativeSpinPTExpt::set_charge_spin(const std::vector& charge_spin) { + assert(inited); + if (settable_chgspin == 0) { + throw deepmd::deepmd_exception( + "this model was not frozen with a charge/spin condition"); + } + if (static_cast(charge_spin.size()) != settable_chgspin) { + throw deepmd::deepmd_exception("the charge/spin condition carries " + + std::to_string(charge_spin.size()) + + " values but the model expects " + + std::to_string(settable_chgspin)); + } + // Route one: the condition of every later forward pass that is not given + // one explicitly. This is the whole mechanism for an uncompressed model, + // which reads the condition as an ordinary input. + default_chg_spin_ = charge_spin; + // Route two: a compressed descriptor has folded the condition into frozen + // tables that the lower holds as constants, so serving another condition + // means rebuilding those tables and writing them over the constants. + if (charge_state_fold_) { + charge_state_fold_->apply(charge_spin, + gpu_enabled ? torch::Device(torch::kCUDA, gpu_id) + : torch::Device(torch::kCPU), + *loader); + } +} + void NativeSpinPTExpt::get_type_map(std::string& type_map_str) { type_map_str.clear(); for (const auto& t : type_map) { @@ -363,7 +444,8 @@ std::vector NativeSpinPTExpt::run_model_graph( const GraphTensorPack& graph, const torch::Tensor& spin, const torch::Tensor& fparam, - const torch::Tensor& aparam) { + const torch::Tensor& aparam, + const torch::Tensor& charge_spin) { deepmd::check_graph_aparam_flat(aparam, daparam, "NativeSpinPTExpt::run_model_graph"); std::vector inputs = { @@ -384,6 +466,9 @@ std::vector NativeSpinPTExpt::run_model_graph( if (daparam > 0) { inputs.push_back(aparam); } + if (dchgspin > 0) { + inputs.push_back(charge_spin); + } return loader->run(inputs); } @@ -412,7 +497,8 @@ std::map NativeSpinPTExpt::run_graph_payload( run_model_graph( graph, spin, make_fparam_tensor(fparam, default_fparam_, dfparam, device), - make_aparam_tensor(aparam, daparam, node_count, nloc, device))); + make_aparam_tensor(aparam, daparam, node_count, nloc, device), + make_chg_spin_tensor(default_chg_spin_, dchgspin, device))); } return output_map; } @@ -956,6 +1042,92 @@ void NativeSpinPTExpt::computew(std::vector& ener, }); } +void NativeSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + check_call_charge_spin(charge_spin, 1, settable_chgspin, + /*applied_per_call=*/false, default_chg_spin_); + computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, fparam, aparam, atomic); +} + +void NativeSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + check_call_charge_spin(charge_spin, 1, settable_chgspin, + /*applied_per_call=*/false, default_chg_spin_); + computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, fparam, aparam, atomic); +} + +void NativeSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + check_call_charge_spin(charge_spin, 1, settable_chgspin, + /*applied_per_call=*/false, default_chg_spin_); + computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, nghost, inlist, ago, fparam, aparam, atomic); +} + +void NativeSpinPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + check_call_charge_spin(charge_spin, 1, settable_chgspin, + /*applied_per_call=*/false, default_chg_spin_); + computew(ener, force, force_mag, virial, atom_energy, atom_virial, coord, + spin, atype, box, nghost, inlist, ago, fparam, aparam, atomic); +} + void NativeSpinPTExpt::compute_canonical_graph_gpu( double* d_atom_energy, double* d_force, diff --git a/source/api_cc/src/commonPTExpt.h b/source/api_cc/src/commonPTExpt.h index dfd6057c59..5ab6a4f42b 100644 --- a/source/api_cc/src/commonPTExpt.h +++ b/source/api_cc/src/commonPTExpt.h @@ -4,17 +4,21 @@ // and helpers for the with-comm dual-artifact layout. #pragma once +#include #include #include #include +#include #include #include #include #include #include +#include #include #include +#include #include #include "common.h" // for remap_comm_sendlist @@ -279,6 +283,67 @@ inline std::vector read_default_chg_spin(const JsonValue& metadata, return default_chg_spin; } +/** + * @brief Validate a charge/spin condition supplied with an inference call. + * + * The condition holds either one frame's values, broadcast to every frame, or + * one set per frame. Whether a call may name a condition of its own depends + * on how the model receives it. A lower that reads the condition as an + * ordinary input takes a different one on every call. A model whose condition + * instead lives in the constants of its compiled tables, and a backend that + * marshals only the condition in force, both serve one state at a time; there + * a supplied condition can only restate the state in force, and anything else + * is rejected rather than silently ignored. + * + * @param[in] charge_spin The condition supplied by the caller. Empty selects + * the state in force and is always accepted. + * @param[in] nframes Number of frames the call evaluates. + * @param[in] settable_chg_spin Width of a charge state the model accepts. + * @param[in] applied_per_call Whether the call marshals the condition into + * the forward pass instead of serving the state in force. + * @param[in] installed The state in force, of width ``settable_chg_spin``. + **/ +inline void check_call_charge_spin(const std::vector& charge_spin, + const int nframes, + const int settable_chg_spin, + const bool applied_per_call, + const std::vector& installed) { + if (charge_spin.empty()) { + return; + } + const std::size_t width = static_cast(settable_chg_spin); + if (charge_spin.size() != width && + charge_spin.size() != width * static_cast(nframes)) { + throw deepmd::deepmd_exception( + "charge_spin has " + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(settable_chg_spin) + " (per frame) or " + + std::to_string(settable_chg_spin * nframes) + " (for " + + std::to_string(nframes) + " frames)."); + } + if (applied_per_call) { + return; + } + if (installed.size() != width) { + throw deepmd::deepmd_exception( + "the model serves one charge/spin state at a time but holds none of " + "the " + + std::to_string(settable_chg_spin) + " values it accepts."); + } + // Charge and multiplicity are integer-valued categorical indices carried as + // double, so a condition either is the state in force or is not. + for (std::size_t ii = 0; ii < charge_spin.size(); ++ii) { + if (charge_spin[ii] != installed[ii % width]) { + throw deepmd::deepmd_exception( + "the charge/spin condition supplied with this call differs from " + "the state the model serves. This model serves one state at a " + "time, held in the constants of its compiled tables or fixed for " + "the whole run, so the state must be chosen with set_charge_spin " + "before inference rather than named per call."); + } + } +} + // ============================================================================ // ZIP archive reader — reads a file from a ZIP archive. // ============================================================================ @@ -509,14 +574,16 @@ class TempFile { * file and return a TempFile owning that path. * * The temp file is created via ``mkstemp(3)`` (atomic, unique, - * 0600 permissions) under the system tempdir (TMPDIR or /tmp). + * 0600 permissions) under the system tempdir (TMPDIR or /tmp), and is + * named after the entry it holds so that a file left behind by a crash + * says which artifact it came from. */ static TempFile from_zip_entry(const std::string& outer_pt2_path, const std::string& entry_name) { std::string content = read_zip_entry(outer_pt2_path, entry_name); const char* tmpdir = std::getenv("TMPDIR"); - std::string tmpl = - std::string(tmpdir ? tmpdir : "/tmp") + "/dp_pt2_with_comm_XXXXXX"; + std::string tmpl = std::string(tmpdir ? tmpdir : "/tmp") + "/dp_pt2_" + + entry_stem(entry_name) + "_XXXXXX"; std::vector buf(tmpl.begin(), tmpl.end()); buf.push_back('\0'); int fd = mkstemp(buf.data()); @@ -548,6 +615,26 @@ class TempFile { } private: + /** + * @brief The base name of a ZIP entry, without directories or extension + * and reduced to characters a file name carries safely. + */ + static std::string entry_stem(const std::string& entry_name) { + const std::size_t slash = entry_name.find_last_of('/'); + std::string stem = + slash == std::string::npos ? entry_name : entry_name.substr(slash + 1); + const std::size_t dot = stem.find_last_of('.'); + if (dot != std::string::npos) { + stem.erase(dot); + } + for (char& c : stem) { + if (!std::isalnum(static_cast(c))) { + c = '_'; + } + } + return stem; + } + void cleanup() { if (!path_.empty()) { ::unlink(path_.c_str()); @@ -658,5 +745,127 @@ inline std::vector build_comm_tensors_positional_with_virtual_atoms( remapped_recvnum.data(), nlocal, nghost); } +// ============================================================================ +// Charge-state fold — the runtime charge/spin condition of a compressed model. +// ============================================================================ + +/** + * @brief The frozen tables through which a compressed descriptor carries its + * charge/spin condition, as a rebuild that can be re-run at any time. + * + * A compressed charge-conditioned descriptor evaluates its frame condition + * once, when the model is frozen, into a handful of tables. Those tables + * reach a compiled lower as module constants, so serving a different + * condition means rebuilding them and writing them over those constants + * rather than re-evaluating the condition on every step. The archive + * therefore ships a second compiled artifact that performs the rebuild, + * together with the name of the constant each of its outputs replaces. + * + * Every lower lifts its constants independently, so the names hold only for + * the lower they were resolved against at freeze time. Only a compressed + * DPA4C descriptor folds a charge state, and that family never carries + * message passing across ranks, so an archive with a fold holds exactly one + * lower and the question of a second set of names does not arise. + * + * An archive without the rebuild leaves the fold inactive. That is the + * ordinary case: an uncompressed model reads its condition as a plain input, + * so nothing needs rebuilding. + */ +class ChargeStateFold { + public: + /** + * @brief Load the rebuild an archive declares, if it declares one. + * + * The constant-name field is the archive's claim that the rebuild ships + * with it, so an archive that declares the names and cannot supply the + * rebuild is malformed and fails here rather than degrading silently. + * + * @param[in] model_path Path to the .pt2 archive. + * @param[in] metadata Parsed archive metadata. + * @param[in] gpu_enabled Whether the lower was loaded on a GPU. + * @param[in] gpu_id The GPU the lower was loaded on. + * @return The fold, or ``nullptr`` when the archive declares none. + **/ + static std::unique_ptr load(const std::string& model_path, + const JsonValue& metadata, + const bool gpu_enabled, + const int gpu_id) { + if (!metadata.obj_val.count("charge_state_constants")) { + return nullptr; + } + if (metadata.obj_val.count("has_comm_artifact") && + metadata["has_comm_artifact"].as_bool()) { + throw deepmd::deepmd_exception( + "the archive ships a charge-state fold beside a with-comm lower; " + "the fold names the constants of one lower only, so the second " + "would keep serving the condition it was frozen against"); + } + std::unique_ptr fold(new ChargeStateFold()); + fold->constants_ = read_names(metadata, "charge_state_constants"); + fold->tempfile_ = std::make_unique( + TempFile::from_zip_entry(model_path, "extra/charge_state.pt2")); + fold->loader_ = std::make_unique( + fold->tempfile_->path(), "model", false, 1, + gpu_enabled ? static_cast(gpu_id) + : static_cast(-1)); + return fold; + } + + /** + * @brief Rebuild the tables for a condition and write them over the + * constants of the lower. + * + * @param[in] charge_spin The condition. + * @param[in] device The device the lower was loaded on. + * @param[in,out] target The lower whose constants carry the condition. + **/ + void apply(const std::vector& charge_spin, + const torch::Device& device, + torch::inductor::AOTIModelPackageLoader& target) const { + // The rebuild consumes the condition in the (1, dim) float32 layout the + // inference lower would receive. + std::vector state(charge_spin.begin(), charge_spin.end()); + torch::Tensor state_tensor = + torch::from_blob(state.data(), + {1, static_cast(state.size())}, + torch::TensorOptions().dtype(torch::kFloat32)) + .clone() + .to(device); + std::vector tables = loader_->run({state_tensor}); + if (tables.size() != constants_.size()) { + throw deepmd::deepmd_exception( + "the charge-state rebuild returned " + std::to_string(tables.size()) + + " tables but the archive names " + std::to_string(constants_.size()) + + " constants; it cannot serve a runtime charge state"); + } + std::unordered_map update; + for (size_t ii = 0; ii < tables.size(); ++ii) { + // An unnamed output belongs to a mechanism this model has disabled and + // has no constant to reach. + if (!constants_[ii].empty()) { + update[constants_[ii]] = tables[ii]; + } + } + target.update_constant_buffer(update, /*use_inactive=*/false, + /*validate_full_updates=*/false); + } + + private: + ChargeStateFold() = default; + + static std::vector read_names(const JsonValue& metadata, + const std::string& key) { + std::vector names; + for (const auto& v : metadata[key].as_array()) { + names.push_back(v.as_string()); + } + return names; + } + + std::vector constants_; + std::unique_ptr tempfile_; + std::unique_ptr loader_; +}; + } // namespace ptexpt } // namespace deepmd diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 06d9f9976a..a03482a90a 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -873,6 +873,22 @@ void PairDeepMD::settings(int narg, char** arg) { "simultaneously"); } + // A charge/spin condition named on the pair_style line holds for the whole + // run, so it is handed to the model once here instead of being resupplied + // every step. This is also what lets a compressed model serve it at all: + // there the condition lives inside frozen tables, which are rebuilt here + // and cannot be rebuilt per step at a sensible cost. + if (!charge_spin.empty()) { + try { + deep_pot.set_charge_spin(charge_spin); + if (numb_models > 1) { + deep_pot_model_devi.set_charge_spin(charge_spin); + } + } catch (deepmd_compat::deepmd_exception& e) { + error->one(FLERR, e.what()); + } + } + if (comm->me == 0) { if (numb_models > 1 && out_freq > 0) { if (!is_restart) { diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index e49c50843a..8ca18f8438 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -201,13 +201,9 @@ void PairDeepMDKokkos::init_style() { "yes' to the input."); } // Runtime frame (fparam) and per-atom (aparam) parameters are threaded to the - // device edge path in compute(); only a runtime charge/spin override is not, - // as compute_edges_gpu draws charge_spin from the model's stored default. - if (!charge_spin.empty()) { - error->all(FLERR, - "pair style deepmd/kk uses the model's stored default " - "charge_spin; a runtime charge_spin is not supported."); - } + // device edge path in compute(). A charge/spin condition needs no threading: + // compute_edges_gpu draws it from the model, and settings() has already fixed + // the model on the condition the pair_style line asked for. // Route the base full request to the Kokkos device neighbor build. auto request = neighbor->find_request(this); diff --git a/source/lmp/pair_deepspin.cpp b/source/lmp/pair_deepspin.cpp index 01d18a944b..d5e3f04efb 100644 --- a/source/lmp/pair_deepspin.cpp +++ b/source/lmp/pair_deepspin.cpp @@ -763,6 +763,22 @@ void PairDeepSpin::settings(int narg, char** arg) { "fparam and fparam_from_compute should NOT be set simultaneously"); } + // A charge/spin condition named on the pair_style line holds for the whole + // run, so it is handed to the model once here instead of being resupplied + // every step. This is also what lets a compressed model serve it at all: + // there the condition lives inside frozen tables, which are rebuilt here + // and cannot be rebuilt per step at a sensible cost. + if (!charge_spin.empty()) { + try { + deep_spin.set_charge_spin(charge_spin); + if (numb_models > 1) { + deep_spin_model_devi.set_charge_spin(charge_spin); + } + } catch (deepmd_compat::deepmd_exception& e) { + error->one(FLERR, e.what()); + } + } + if (comm->me == 0) { if (numb_models > 1 && out_freq > 0) { if (!is_restart) { diff --git a/source/lmp/pair_dpa4spin.cpp b/source/lmp/pair_dpa4spin.cpp index b68b7622f3..501a49c869 100644 --- a/source/lmp/pair_dpa4spin.cpp +++ b/source/lmp/pair_dpa4spin.cpp @@ -224,10 +224,10 @@ void PairDPA4Spin::settings(int narg, char** arg) { // Name whichever style the input selected, so the Kokkos variant reports // itself rather than its host base. const std::string style = force->pair_style; - if (narg != 1) { + if (narg < 1) { error->all(FLERR, "Illegal pair_style command: pair style " + style + " evaluates a single native-spin artifact and takes " - "its path as the only argument."); + "its path as the first argument."); } try { @@ -237,12 +237,74 @@ void PairDPA4Spin::settings(int narg, char** arg) { } cutoff = deep_spin.cutoff() * dist_unit_cvt_factor; + // How many values name a charge state is a property of the artifact, so the + // keyword can only be read once the model is loaded. + const int dim_chg_spin = deep_spin.dim_chg_spin(); + std::vector charge_spin; + int iarg = 1; + while (iarg < narg) { + if (std::string(arg[iarg]) != "charge_spin") { + error->all(FLERR, "Illegal pair_style command: pair style " + style + + " takes the artifact path followed by the optional " + "keyword charge_spin, not '" + + std::string(arg[iarg]) + "'."); + } + // One charge state holds for the whole run, so a second occurrence names + // no state the style could serve. Values already read mark the keyword + // as seen: a successful read always contributes at least one value. + if (!charge_spin.empty()) { + error->all(FLERR, + "Illegal pair_style command: keyword charge_spin names the " + "single charge state of the whole run, so it may appear " + "only once."); + } + if (dim_chg_spin == 0) { + error->all(FLERR, + "Illegal pair_style command: the artifact served by pair " + "style " + + style + + " carries no charge/spin condition, so keyword " + "charge_spin names nothing it can serve."); + } + if (iarg + dim_chg_spin >= narg) { + error->all(FLERR, + "Illegal pair_style command: keyword charge_spin names a " + "charge state with " + + std::to_string(dim_chg_spin) + " value(s)."); + } + for (int ii = 0; ii < dim_chg_spin; ++ii) { + charge_spin.push_back( + utils::numeric(FLERR, arg[iarg + 1 + ii], false, lmp)); + } + iarg += 1 + dim_chg_spin; + } + + // A charge/spin condition named on the pair_style line holds for the whole + // run, so it is handed to the model once here instead of being resupplied + // every step. This is also what lets a compressed model serve it at all: + // there the condition lives inside frozen tables, which are rebuilt here + // and cannot be rebuilt per step at a sensible cost. + if (!charge_spin.empty()) { + try { + deep_spin.set_charge_spin(charge_spin); + } catch (deepmd_compat::deepmd_exception& e) { + error->one(FLERR, e.what()); + } + } + utils::logmesg(lmp, " >>> Info of model(s):\n" " using 1 model(s): {}\n" " rcut in model: {}\n" " ntypes in model: {}\n", arg[0], cutoff, deep_spin.numb_types()); + if (!charge_spin.empty()) { + std::string values; + for (const double value : charge_spin) { + values += fmt::format("{} ", value); + } + utils::logmesg(lmp, " using charge_spin: {}\n", values); + } } /* ---------------------------------------------------------------------- diff --git a/source/lmp/pair_dpa4spin.h b/source/lmp/pair_dpa4spin.h index 73ac544734..7d0aba8977 100644 --- a/source/lmp/pair_dpa4spin.h +++ b/source/lmp/pair_dpa4spin.h @@ -54,8 +54,10 @@ class CommBrickDPA4Spin : public CommBrick { // ghost) force and magnetic force, both of which the spin atom style folds // onto their owners through its reverse communication. // -// The style evaluates exactly one artifact and passes no frame, atomic or -// charge/spin parameters, so a model that requires one must carry its default. +// The style evaluates exactly one artifact and passes no frame or atomic +// parameters, so a model that requires one must carry its default. A charge +// state, being fixed for the whole run, is named once on the pair_style line +// and handed to the model there. // // The device-resident variant is ``dpa4spin/kk``; it needs the compact // canonical artifact and evaluates it without the per-step host marshaling. @@ -64,7 +66,8 @@ class PairDPA4Spin : public Pair { PairDPA4Spin(class LAMMPS*); ~PairDPA4Spin() override; - // Load the artifact named by ``pair_style dpa4spin ``. + // Load the artifact named by ``pair_style dpa4spin [charge_spin + // ... ]`` and fix the charge state the keyword names. void settings(int, char**) override; // Resolve the LAMMPS atom types onto the element list of the model. void coeff(int, char**) override; diff --git a/source/tests/common/dpmodel/test_descriptor_dpa4c.py b/source/tests/common/dpmodel/test_descriptor_dpa4c.py index 46ab67a072..4c7c583cfd 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa4c.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa4c.py @@ -106,7 +106,7 @@ def edge_features( return descriptor.build_edge_features( graph, atype_local, - descriptor.pair_film.call(descriptor.type_embedding.call()), + descriptor.pair_film.pair_latent(descriptor.type_embedding.call()), )[:3] @@ -971,7 +971,10 @@ def moment_divisor(self, spin: np.ndarray) -> np.ndarray: *self.descriptor.build_edge_features( self.graph, self.atype, - self.descriptor.pair_film.call(self.descriptor.type_embedding.call()), + self.descriptor.pair_film.pair_latent( + self.descriptor.type_embedding.call() + ), + None, masked, ), self.descriptor.spin.onsite_payload(masked, self.atype), @@ -1361,3 +1364,343 @@ def test_calibration_rejects_a_corpus_without_moments() -> None: corpus[0].pop("spin") with pytest.raises(ValueError, match="requires a per-node magnetic"): calibrated_descriptor(corpus) + + +#: Frame conditions exercised by the charge-state tests, as +#: ``[charge, multiplicity]`` pairs. +NEUTRAL_SINGLET = np.array([[0.0, 1.0]]) +CATION_TRIPLET = np.array([[2.0, 3.0]]) + + +def make_charge_descriptor( + *, + seed: int = 17, + activate: bool = True, + **kwargs: Any, +) -> DescrptDPA4C: + """Return a charge-conditioned descriptor. + + Parameters + ---------- + seed + Parameter-initialization seed. + activate + Whether to replace the zero-initialized condition output head with + deterministic weights. Every property that depends on the condition + actually reaching the descriptor needs an active head; leaving it at + its initialization is itself the subject of one test. + **kwargs + Overrides forwarded to the descriptor constructor. + + Returns + ------- + DescrptDPA4C + Charge-conditioned descriptor. + """ + kwargs.setdefault("default_chg_spin", [0.0, 1.0]) + descriptor = DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=8, + lmax=2, + n_radial=4, + precision="float64", + seed=seed, + add_chg_spin_ebd=True, + **kwargs, + ) + if activate: + head = descriptor.charge_spin_embedding.network.layers[-1] + head.w = np.random.default_rng(seed).normal(0.0, 0.5, size=head.w.shape) + return descriptor + + +def evaluate_conditioned( + descriptor: DescrptDPA4C, + charge_spin: np.ndarray | None, + coord: np.ndarray = COORD, + atype: np.ndarray = ATYPE, +) -> np.ndarray: + """Evaluate a charge-conditioned descriptor on the graph interface.""" + graph, atype_local = build_graph(descriptor, coord, atype) + return descriptor.call_graph(graph, atype_local, charge_spin=charge_spin)[0] + + +def charge_route_heads(descriptor: DescrptDPA4C) -> tuple[np.ndarray, np.ndarray]: + """Split the condition output head into its two routes. + + Returns + ------- + type_route + Head restricted to the centre type-embedding columns. + pair_route + Head restricted to the ordered pair encoder columns. + """ + weight = descriptor.charge_spin_embedding.network.layers[-1].w + type_route, pair_route = weight.copy(), weight.copy() + type_route[:, descriptor.channels :] = 0.0 + pair_route[:, : descriptor.channels] = 0.0 + return type_route, pair_route + + +def test_an_unconditioned_descriptor_declares_no_frame_condition() -> None: + descriptor = DescrptDPA4C(rcut=3.0, ntypes=2, channels=8, lmax=2, n_radial=4) + assert descriptor.charge_spin_embedding is None + assert not descriptor.supports_charge_spin() + assert descriptor.get_dim_chg_spin() == 0 + assert not descriptor.has_default_chg_spin() + + +def test_an_untrained_condition_head_reproduces_the_plain_descriptor() -> None: + """The condition output projection starts at zero. + + An untrained descriptor is therefore independent of the charge state for + every value of it, so the fixed output calibration measured once before + training carries no random condition offset. + """ + plain = DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=8, + lmax=2, + n_radial=4, + precision="float64", + seed=17, + ) + conditioned = make_charge_descriptor(activate=False) + reference = evaluate(plain).reshape(-1, plain.get_dim_out()) + for condition in (NEUTRAL_SINGLET, CATION_TRIPLET): + np.testing.assert_array_equal( + evaluate_conditioned(conditioned, condition), + reference, + ) + + +def test_each_condition_route_reaches_the_descriptor() -> None: + """Both injection points must be live. + + The centre type route alone would be indistinguishable from handing the + condition to the fitting network as a frame parameter. Only the ordered + pair route changes how a given geometry maps to the degree-wise moments, + so the two are asserted separately rather than through their sum. + """ + descriptor = make_charge_descriptor() + head = descriptor.charge_spin_embedding.network.layers[-1] + for route in charge_route_heads(descriptor): + head.w = route + assert not np.allclose( + evaluate_conditioned(descriptor, NEUTRAL_SINGLET), + evaluate_conditioned(descriptor, CATION_TRIPLET), + ) + + +def test_frames_carry_independent_conditions() -> None: + """A batched evaluation must agree with per-frame evaluations. + + Each frame occupies one contiguous block of the flat node axis and an + edge inherits the frame of the centre it reduces onto, so a batch of + mixed charge states is only correct if that map is exact. + """ + descriptor = make_charge_descriptor() + shifted = COORD + 0.05 + batched = evaluate_conditioned( + descriptor, + np.concatenate([NEUTRAL_SINGLET, CATION_TRIPLET], axis=0), + np.concatenate([COORD, shifted], axis=0), + np.concatenate([ATYPE, ATYPE], axis=0), + ) + np.testing.assert_allclose( + batched, + np.concatenate( + [ + evaluate_conditioned(descriptor, NEUTRAL_SINGLET, COORD), + evaluate_conditioned(descriptor, CATION_TRIPLET, shifted), + ], + axis=0, + ), + atol=1e-12, + ) + + +def test_a_frame_condition_preserves_rotation_invariance() -> None: + # The condition is a pair of scalars, so it may not disturb the O(3) + # invariance the readout establishes. + descriptor = make_charge_descriptor() + angle = 0.7 + rotation = np.array( + [ + [math.cos(angle), -math.sin(angle), 0.0], + [math.sin(angle), math.cos(angle), 0.0], + [0.0, 0.0, 1.0], + ] + ) + np.testing.assert_allclose( + evaluate_conditioned(descriptor, CATION_TRIPLET, COORD @ rotation.T), + evaluate_conditioned(descriptor, CATION_TRIPLET), + atol=1e-12, + ) + + +def test_a_missing_condition_falls_back_to_the_configured_default() -> None: + descriptor = make_charge_descriptor() + np.testing.assert_array_equal( + evaluate_conditioned(descriptor, None), + evaluate_conditioned(descriptor, NEUTRAL_SINGLET), + ) + + +def test_a_missing_condition_without_a_default_is_an_error() -> None: + descriptor = make_charge_descriptor(default_chg_spin=None) + with pytest.raises(ValueError, match="requires a frame `charge_spin`"): + evaluate_conditioned(descriptor, None) + + +def test_the_dense_interface_conditions_on_the_frame_state() -> None: + # The dense adapter is the reference path of the common descriptor ABI; + # dropping the condition there would be silent. + descriptor = make_charge_descriptor() + coord_ext, atype_ext, mapping, nlist = dense_inputs(descriptor) + dense = descriptor( + coord_ext, + atype_ext, + nlist, + mapping=mapping, + charge_spin=CATION_TRIPLET, + )[0] + np.testing.assert_allclose( + dense.reshape(-1, descriptor.get_dim_out()), + evaluate_conditioned(descriptor, CATION_TRIPLET), + atol=1e-12, + ) + + +def test_serialization_preserves_the_conditioned_descriptor() -> None: + descriptor = make_charge_descriptor() + restored = DescrptDPA4C.deserialize(descriptor.serialize()) + assert restored.get_default_chg_spin() == descriptor.get_default_chg_spin() + np.testing.assert_allclose( + evaluate_conditioned(restored, CATION_TRIPLET), + evaluate_conditioned(descriptor, CATION_TRIPLET), + atol=1e-14, + ) + + +def test_the_frozen_pair_cache_reproduces_the_per_edge_route() -> None: + """Compression folds the condition into the ordered pair cache. + + Training evaluates the conditioning heads on the edge axis, because the + product of the frame and ordered-pair axes exceeds the edge count for the + molecular systems a charge state describes. Compression evaluates the + same heads once over the finite pair table. The compressed artifact is + only valid because the two agree exactly. + """ + descriptor = make_charge_descriptor() + type_embedding = descriptor.type_embedding.call() + _type_shift, pair_hidden_bias = descriptor.charge_spin_embedding.call( + CATION_TRIPLET + ) + folded = descriptor.pair_film.call(type_embedding, hidden_bias=pair_hidden_bias[0]) + pre_activation, base_shift = descriptor.pair_film.pair_latent(type_embedding) + pair_index = np.arange(pre_activation.shape[0]) + per_edge = descriptor.build_pair_conditioning( + (pre_activation, base_shift), + pair_index, + np.broadcast_to( + pair_hidden_bias, (pair_index.size, pair_hidden_bias.shape[-1]) + ), + ) + for cache, edge in zip(folded, per_edge, strict=True): + if cache is None: + assert edge is None + else: + np.testing.assert_allclose(edge, cache, atol=1e-14) + + +def test_the_padding_type_keeps_its_zero_centre_features() -> None: + """The condition shifts only the real rows of the centre type table. + + Compressed inference conditions a frozen table whose padding row stays + zero, so shifting that row on the portable path would break the parity + between the two. + """ + descriptor = make_charge_descriptor() + type_embedding = descriptor.type_embedding.call() + atype = np.array([0, 1, descriptor.ntypes], dtype=np.int64) + type_shift, _pair_hidden_bias = descriptor.charge_spin_embedding.call( + CATION_TRIPLET + ) + features = descriptor.build_center_type_features( + type_embedding, + atype, + np.broadcast_to(type_shift, (atype.size, descriptor.channels)), + ) + np.testing.assert_array_equal(features[2], np.zeros(descriptor.channels)) + assert not np.allclose(features[0], type_embedding[0]) + + +def charge_corpus(charge_spin: np.ndarray | None) -> list[dict]: + """Build a two-type calibration corpus carrying one frame condition.""" + rng = np.random.default_rng(3) + nframes, natoms, cell = 6, 8, 9.0 + system = { + "coord": rng.uniform(0.0, cell, size=(nframes, natoms, 3)), + "atype": np.tile(np.array([0, 1, 0, 1, 0, 1, 0, 1]), (nframes, 1)), + "box": np.tile(np.diag([cell, cell, cell]).reshape(1, 9), (nframes, 1)), + } + if charge_spin is not None: + system["charge_spin"] = charge_spin + return [system] + + +@pytest.mark.parametrize( + "charge_spin", + [ + np.tile(np.array([[-1.0, 2.0]]), (6, 1)), + np.array([[-1.0, 2.0]]), + np.array([-1.0, 2.0]), + ], + ids=["per-frame", "single-pair-2d", "single-pair-1d"], +) +def test_the_calibration_accepts_every_shape_evaluation_accepts( + charge_spin: np.ndarray, +) -> None: + """A system may state one condition for all of its frames. + + Evaluation broadcasts a single pair over the frame axis, so a calibration + that required one row per frame would reject a corpus the trained model + then runs on without complaint. + """ + descriptor = make_charge_descriptor() + descriptor.compute_input_stats(charge_corpus(charge_spin)) + frames = descriptor._calibration_frames(charge_corpus(charge_spin)[0]) + for frame in frames: + np.testing.assert_array_equal(frame["charge_spin"], np.array([[-1.0, 2.0]])) + + +def test_the_calibration_samples_the_corpus_charge_states() -> None: + """The preconditioner must be measured over the sampled charge states. + + It is frozen once and has to hold for every state the corpus contains, so + a calibration that read one state, or none, would fix it on the wrong + scale. + """ + descriptor = make_charge_descriptor() + rng = np.random.default_rng(3) + nframes, natoms, cell = 6, 8, 9.0 + corpus = [ + { + "coord": rng.uniform(0.0, cell, size=(nframes, natoms, 3)), + "atype": np.tile(np.array([0, 1, 0, 1, 0, 1, 0, 1]), (nframes, 1)), + "box": np.tile(np.diag([cell, cell, cell]).reshape(1, 9), (nframes, 1)), + "charge_spin": np.tile(np.array([[-1.0, 2.0]]), (nframes, 1)), + } + ] + frames = descriptor._calibration_frames(corpus[0]) + np.testing.assert_array_equal(frames[0]["charge_spin"], np.array([[-1.0, 2.0]])) + + descriptor.compute_input_stats(corpus) + default_state = make_charge_descriptor() + default_state.compute_input_stats( + [{key: value for key, value in corpus[0].items() if key != "charge_spin"}] + ) + assert not np.allclose(descriptor.stddev, default_state.stddev) diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py index e9f2bd3973..9e3f41a24a 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -365,7 +365,7 @@ def test_amp_spans_the_edge_stage_and_stops_at_its_boundary( features = descriptor.build_edge_features( graph, atype_local, - descriptor.pair_film.call(descriptor.type_embedding.call()), + descriptor.pair_film.pair_latent(descriptor.type_embedding.call()), )[:3] finally: for handle in handles: @@ -456,7 +456,9 @@ def setup_method(self) -> None: self.spin = torch.randn(6, 3, dtype=torch.float64, generator=generator).to( env.DEVICE ) - from deepmd.dpmodel.utils.neighbor_graph import build_neighbor_graph + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) self.graph = build_neighbor_graph(self.coord, self.atype, None, 3.0) self.flat_atype = self.atype.reshape(-1) @@ -551,7 +553,9 @@ def test_compression_covers_the_spin_families(self) -> None: structural set as a spin-free one and its frozen tables are built alongside the geometric caches. """ - from deepmd.kernels.cuda.dpa4c.graph_compress import mega_eligible + from deepmd.kernels.cuda.dpa4c.graph_compress import ( + mega_eligible, + ) single = DescrptDPA4C( rcut=3.0, diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py index f710d8301c..2ab071467d 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py @@ -406,6 +406,119 @@ def test_post_compression_statistics_update_snapshot( torch.testing.assert_close(actual, reference, atol=3e-5, rtol=3e-5) +def _build_charge_descriptor( + channels: int = 8, + radial_modes: int = 0, + default_chg_spin: list[float] | None = [2.0, 3.0], +) -> DescrptDPA4C: + """Return a charge-conditioned descriptor with an active condition head. + + The condition output projection is zero initialized, so it is replaced by + deterministic weights: a parity test against an inert condition would + pass for the wrong reason. + """ + descriptor = ( + DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=channels, + lmax=2, + n_radial=8, + radial_modes=radial_modes, + precision="float32", + seed=17, + add_chg_spin_ebd=True, + default_chg_spin=default_chg_spin, + ) + .cuda() + .eval() + ) + head = descriptor.charge_spin_embedding.network.layers[-1] + generator = torch.Generator(device="cuda").manual_seed(11) + with torch.no_grad(): + head.w.copy_( + torch.randn( + head.w.shape, + dtype=torch.float32, + device="cuda", + generator=generator, + ) + * 0.5 + ) + return descriptor + + +@_GPU +@pytest.mark.parametrize("channels", [8, 32, 128]) +@pytest.mark.parametrize("radial_modes", [0, 4]) +def test_charge_condition_folds_into_the_frozen_tables( + monkeypatch: pytest.MonkeyPatch, + channels: int, + radial_modes: int, +) -> None: + """Compression must reproduce the conditioned portable descriptor. + + The frame condition reaches only the finite type and ordered pair tables, + so folding it there leaves every artifact shape and the compiled kernel + untouched. Parity against the portable path is what establishes that the + fold is the same function the edge axis evaluates. + """ + descriptor = _build_charge_descriptor(channels, radial_modes) + graph, atype = _build_graph(descriptor, canonical=False) + reference, _ = descriptor.call_graph( + graph, + atype, + charge_spin=torch.tensor([[2.0, 3.0]], device="cuda"), + ) + descriptor.enable_compression(min_nbor_dist=0.5) + monkeypatch.setenv("DP_CUDA_INFER", "1") + actual, _ = descriptor.call_graph(graph, atype) + torch.testing.assert_close(actual, reference, atol=3e-5, rtol=3e-5) + + +@_GPU +def test_a_baked_charge_state_differs_from_a_neutral_one() -> None: + """The baked state must actually reach the frozen tables. + + Two snapshots of the same weights that differ only in the charge state + they were compressed against have to disagree; otherwise the fold is + writing an unconditioned table and the parity test above would hold + vacuously. + """ + graph, atype = _build_graph(_build_charge_descriptor(), canonical=False) + outputs = [] + for state in ([0.0, 1.0], [2.0, 3.0]): + descriptor = _build_charge_descriptor(default_chg_spin=state) + descriptor.enable_compression(min_nbor_dist=0.5) + outputs.append( + torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec, + *_arguments(descriptor, graph, atype), + )[0] + ) + assert not torch.allclose(outputs[0], outputs[1]) + + +@_GPU +def test_a_compressed_snapshot_declares_no_runtime_condition() -> None: + # The compact canonical lower carries no conditioning slot, so a baked + # snapshot has to report that it consumes none. + descriptor = _build_charge_descriptor() + assert descriptor.get_dim_chg_spin() == 2 + descriptor.enable_compression(min_nbor_dist=0.5) + assert descriptor.get_dim_chg_spin() == 0 + assert descriptor.supports_charge_spin() + + +@_GPU +def test_compression_requires_a_baked_charge_state() -> None: + # Without a default there is no state to fold, and a snapshot that + # silently evaluated the unconditioned tables would be a different model. + descriptor = _build_charge_descriptor(default_chg_spin=None) + with pytest.raises(ValueError, match="`default_chg_spin`"): + descriptor.enable_compression(min_nbor_dist=0.5) + + @_GPU def test_compressed_descriptor_cannot_reenter_training() -> None: descriptor = _build_descriptor(8) diff --git a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py index 376b679c31..7ed04e249d 100644 --- a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py @@ -184,6 +184,117 @@ def test_auto_lower_kind_selects_compact_canonical(channels: int) -> None: assert _resolve_lower_kind("model.pt2", data, "auto") == "dpa4c_canonical" +def test_the_graph_lower_conditions_each_frame_on_its_own_charge_state() -> None: + """The frame condition must survive the model seam and stay per frame. + + The atomic model forwards ``charge_spin`` only to descriptors that + declare the capability, and the graph lower flattens every frame onto one + node axis, so a batch of mixed charge states exercises both the seam and + the node-to-frame map. + """ + config = _config() + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [0.0, 1.0] + model = get_model(config).to(env.DEVICE).eval() + descriptor = model.get_descriptor() + assert descriptor.supports_charge_spin() + + # The condition output projection is zero initialized, so an untrained + # model would be inert and the comparison below would hold vacuously. + head = descriptor.charge_spin_embedding.network.layers[-1] + generator = torch.Generator(device=head.w.device).manual_seed(5) + with torch.no_grad(): + head.w.copy_( + torch.randn( + head.w.shape, + dtype=head.w.dtype, + device=head.w.device, + generator=generator, + ) + * 0.5 + ) + + sample = build_synthetic_graph_inputs( + model, + e_max=None, + nframes=2, + nloc=7, + dtype=torch.float64, + device=env.DEVICE, + ) + + def energy(second_state: list[float]) -> torch.Tensor: + return model.forward_common_lower_graph( + *sample[:10], + destination_sorted=True, + do_atomic_virial=True, + fparam=sample[10], + aparam=sample[11], + charge_spin=torch.tensor( + [[0.0, 1.0], second_state], + dtype=torch.float64, + device=env.DEVICE, + ), + )["energy_redu"] + + neutral, mixed = energy([0.0, 1.0]), energy([2.0, 3.0]) + torch.testing.assert_close(neutral[0], mixed[0], atol=0.0, rtol=0.0) + assert float((mixed[1] - neutral[1]).detach().abs().max()) > 0.0 + + +def test_an_uncompressed_export_keeps_the_charge_state_as_a_runtime_input() -> None: + """An uncompressed artifact conditions at run time, not at export time. + + Whether a deployed model accepts a charge state is a property of + compression rather than of the export format: the graph lower carries a + conditioning slot with a dynamic frame axis, and only the fold of the + compact canonical path removes it. + """ + config = _compressed_config() + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [2.0, 3.0] + model = get_model(config).to("cpu").eval() + exported, metadata, _model_json, _output_keys = _trace_and_export( + {"model": model.serialize()}, + lower_kind="graph", + do_atomic_virial=True, + ) + assert metadata["dim_chg_spin"] == 2 + assert metadata["default_chg_spin"] == [2.0, 3.0] + placeholders = [ + node.name + for node in exported.graph_module.graph.nodes + if node.op == "placeholder" + ] + assert placeholders[-1].startswith("charge_spin") + + +def test_a_baked_charge_state_reaches_the_compact_canonical_lower() -> None: + """Compression must remove the runtime condition, not just satisfy it. + + The compact canonical argument list carries no conditioning slot, and + evaluation rejects an artifact that claims to need one. A charge- + conditioned model reaches that lower only because compression folds the + charge state into the frozen tables and the snapshot then reports a zero + runtime condition width. + """ + config = _compressed_config() + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [2.0, 3.0] + model = get_model(config).to("cpu").eval() + assert model.get_dim_chg_spin() == 2 + assert _resolve_lower_kind("model.pt2", {"model": model.serialize()}, "auto") == ( + "graph" + ) + + descriptor = model.get_descriptor() + descriptor.enable_compression(min_nbor_dist=0.5) + assert model.get_dim_chg_spin() == 0 + assert _resolve_lower_kind("model.pt2", {"model": model.serialize()}, "auto") == ( + "dpa4c_canonical" + ) + + def test_compact_canonical_eligibility_rejects_other_descriptors() -> None: from deepmd.kernels.cuda.dpa4c.canonical import ( canonical_model_eligible, From 99c84681dec416aa01d40c40af9f099ea6721d60 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 9 Aug 2026 00:41:56 +0800 Subject: [PATCH 04/10] feat(dpa4c): fine-tune native spin from a spin-free pretraining Naming a magnetic type on a pretraining that declared none must leave the predicted energy untouched, because the spin routes the activation releases never received a gradient. On FeC it otherwise moves the energy by several eV per atom with a configuration-dependent sign, which is what forces the output-bias regression to solve for a per-type constant of several keV. A single scalar gate on the whole spin branch makes the activation exact. It multiplies the block after the calibration, so a closed gate feeds the fitting network exactly zero whatever preconditioner was measured, and the invariants are linear in it, so zero is a starting point whose gradient is the branch itself rather than a stationary point. No weight inside the branch can play that role: the families reach the fitting network by several routes and at two spin orders, and a factor on the conditioned moment would enter the degree-one Grams squared and the quadrupole Grams to the fourth power. Constructing the gate closed is the whole mechanism. A transfer either copies a closed gate or keeps the freshly constructed one, so no reset hook takes part, and a checkpoint predating the gate carries no value for it that the runtime would invent. Fine-tuning such a corpus needs batches of unequal atom count. The graph lower already reads a flat node axis, so a ragged batch feeds it directly while a rectangular one has its phantom padding compacted away before the network sees it, and the loss reads each frame's own atom count from the graph. Two corrections ride along. Calibration accepts every scale the storage precision can represent instead of rejecting representable extremes, and the node-backward group width follows the occupancy the running device reports rather than a compiled-in constant, which the launch bounds of newer architectures ignore. --- deepmd/dpmodel/descriptor/dpa4c.py | 93 +++++-- deepmd/dpmodel/descriptor/dpa4c_nn/spin.py | 23 +- deepmd/kernels/cuda/dpa4c/graph_compress.py | 31 +++ deepmd/pt_expt/descriptor/dpa4c.py | 1 + deepmd/pt_expt/model/ener_model.py | 124 +++++++--- deepmd/pt_expt/train/training.py | 173 ++++++------- source/op/pt/dpa4c_graph_compress_kernel.cuh | 68 ++++- source/op/pt/dpa4c_graph_compress_launch.h | 21 +- .../common/dpmodel/test_descriptor_dpa4c.py | 98 +++++++- source/tests/pt_expt/descriptor/test_dpa4c.py | 65 ++++- .../pt_expt/descriptor/test_dpa4c_cuda.py | 6 +- .../pt_expt/model/test_dpa4c_graph_lower.py | 233 ++++++++++++++++++ source/tests/pt_expt/test_finetune.py | 12 + 13 files changed, 784 insertions(+), 164 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index f413adf126..d714a5f515 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -1144,7 +1144,28 @@ def build_invariant_descriptor( descriptor = xp.concat(blocks, axis=-1) mean = xp_asarray_nodetach(xp, self.mean, device=device) stddev = xp_asarray_nodetach(xp, self.stddev, device=device) - return (descriptor - mean[None, :]) / stddev[None, :] + calibrated = (descriptor - mean[None, :]) / stddev[None, :] + if self.spin is None: + return calibrated + # === Spin branch gate === + # Applied to the CALIBRATED block, so a closed gate feeds the fitting + # network exactly zero whatever calibration was measured, and the + # compressed path reproduces it by scaling the inverse deviation of + # these columns alone -- an affine map that stays exact at zero. + gate = xp.astype( + xp_asarray_nodetach(xp, self.spin.spin_gate[...], device=device), + calibrated.dtype, + ) + start = self.readout.get_dim_out() + stop = start + self.spin.get_dim_out() + return xp.concat( + [ + calibrated[:, :start], + calibrated[:, start:stop] * gate, + calibrated[:, stop:], + ], + axis=-1, + ) # === Backend primitives === # A backend wrapper overrides these to reach native kernels; the equations @@ -1491,29 +1512,31 @@ def compute_input_stats( geometric = np.zeros(geometry_dim, dtype=bool) geometric[: self.readout.get_dim_out()] = True geometric[geometry_dim - 2 :] = True - degenerate = geometric & ( - ~np.isfinite(feature_rms) | (feature_rms <= self._STAT_EPS) + finite_positive = np.isfinite(feature_rms) & (feature_rms > 0.0) + invalid = measured & ( + ~np.isfinite(feature_rms) | (geometric & ~finite_positive) ) - if np.any(degenerate): + if np.any(invalid): raise ValueError( - "DPA4C output calibration requires non-degenerate finite " - f"geometric features, got RMS values {feature_rms.tolist()}" + "DPA4C output calibration requires finite measured feature " + "RMS values and strictly positive geometric RMS values, got " + "invalid output indices " + f"{np.array2string(np.flatnonzero(invalid), threshold=16)}" ) type_table = to_numpy_array(self.type_embedding.call())[: self.ntypes] target_rms = float(np.sqrt(np.mean(np.square(type_table, dtype=np.float64)))) - if not math.isfinite(target_rms) or target_rms <= self._STAT_EPS: + if not math.isfinite(target_rms) or target_rms <= 0.0: raise ValueError( f"DPA4C type embedding has a degenerate calibration RMS {target_rms}" ) - # A coordinate earns a preconditioner only where its measured scale is - # meaningful. The geometric block is already required to be - # non-degenerate above; the spin block is not, because a corpus whose - # moments are uniformly weak drives the quartic spin coordinates to a - # vanishing root mean square, and dividing by it would hand them an - # unbounded gain. Those coordinates keep the identity instead, which is - # the same treatment a coordinate that never activates receives. - conditioned = measured & np.isfinite(feature_rms) - conditioned &= feature_rms > self._STAT_EPS + # Geometric polynomial families span different degrees and therefore + # have no common absolute RMS threshold. Every positive finite + # geometric coordinate is calibrated. Weak spin coordinates retain + # the identity scale because their sampled magnitude can reflect the + # magnetic population rather than geometric degeneracy. The complete + # preconditioner is validated in storage precision below. + conditioned = measured & finite_positive + conditioned &= geometric | (feature_rms > self._STAT_EPS) geometry_stddev = np.ones(geometry_dim, dtype=np.float64) geometry_stddev[conditioned] = feature_rms[conditioned] / target_rms geometry_mean = np.zeros(geometry_dim, dtype=np.float64) @@ -1539,12 +1562,40 @@ def compute_input_stats( geometry_stddev[mass] = mass_stddev / target_rms tail = np.zeros(self.channels, dtype=np.float64) - self.mean = np.concatenate([geometry_mean, tail]).astype( - PRECISION_DICT[self.precision] - ) - self.stddev = np.concatenate([geometry_stddev, tail + 1.0]).astype( - PRECISION_DICT[self.precision] + output_mean = np.concatenate([geometry_mean, tail]) + output_stddev = np.concatenate([geometry_stddev, tail + 1.0]) + storage_dtype = np.dtype(PRECISION_DICT[self.precision]) + with np.errstate( + divide="ignore", + invalid="ignore", + over="ignore", + under="ignore", + ): + stored_mean = output_mean.astype(storage_dtype) + stored_stddev = output_stddev.astype(storage_dtype) + inverse_stddev = np.reciprocal(stored_stddev) + invalid_mean = ~np.isfinite(stored_mean) + invalid_scale = ( + ~np.isfinite(stored_stddev) + | (stored_stddev <= 0.0) + | ~np.isfinite(inverse_stddev) + | (inverse_stddev <= 0.0) ) + if np.any(invalid_mean): + raise ValueError( + "DPA4C output calibration produced non-representable means " + f"in {storage_dtype.name} at output indices " + f"{np.array2string(np.flatnonzero(invalid_mean), threshold=16)}" + ) + if np.any(invalid_scale): + raise ValueError( + "DPA4C output calibration produced scales that are not " + f"positive, finite, and invertible in {storage_dtype.name} " + "at output indices " + f"{np.array2string(np.flatnonzero(invalid_scale), threshold=16)}" + ) + self.mean = stored_mean + self.stddev = stored_stddev def _calibration_frames(self, system: dict) -> list[dict]: """Draw the calibration frames of one sampled system. diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py index d20075571c..56ea73389f 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py @@ -211,6 +211,8 @@ class SpinChannels(NativeOP): If ``use_spin`` does not have one entry per real atom type. """ + CONFIG_DERIVED_ARRAYS = ("spin_mask",) + def __init__( self, ntypes: int, @@ -234,8 +236,9 @@ def __init__( precision_dtype = PRECISION_DICT[self.precision.lower()] # === Per-type spin gate === - # Deterministic from the configuration, so it is rebuilt rather than - # serialized. The trailing row is the padding type. + # Deterministic from the configuration (hence ``CONFIG_DERIVED_ARRAYS``), + # so it is rebuilt rather than serialized. The trailing row is the + # padding type. self.spin_mask = np.asarray( [1.0 if flag else 0.0 for flag in self.use_spin] + [0.0], dtype=precision_dtype, @@ -258,6 +261,18 @@ def __init__( 0.0, 1.0, size=(self.ntypes + 1,) ).astype(precision_dtype) + # === Branch gate === + # One scalar on the whole branch, the only place all of it passes + # through: the families reach the fitting network by several routes + # and at two spin orders, so no weight inside the branch gates all of + # them, and a factor applied to the conditioned moment instead would + # enter the invariants quadratically and leave zero a stationary + # point. The descriptor applies it to the CALIBRATED block (see + # ``DescrptDPA4C._readout``), so a closed gate contributes exactly + # zero whatever the measured calibration is, and the compressed path + # carries it as a factor on the inverse deviation alone. + self.spin_gate = np.zeros((1,), dtype=precision_dtype) + # Isometric half-vectorization of the two spin Grams. The quadrupole # block drops entry zero, its on-site self-term: the harmonic blocks # are homogeneous, so |B_2(s)|^2 = |s|^4 exactly and that entry is a @@ -651,6 +666,7 @@ def serialize(self) -> dict[str, Any]: "adam_spin_quadrupole_weight": to_numpy_array( self.adam_spin_quadrupole_weight ), + "spin_gate": to_numpy_array(self.spin_gate), }, } @@ -677,7 +693,7 @@ def deserialize(cls, data: dict[str, Any]) -> SpinChannels: check_version_compatibility(data.pop("@version"), 1, 1) if data.pop("@class") != "SpinChannels": raise ValueError("Invalid serialized class for SpinChannels") - variables = data.pop("@variables") + variables = dict(data.pop("@variables")) obj = cls(**data) obj.set_variables(variables) return obj @@ -700,6 +716,7 @@ def set_variables(self, variables: dict[str, Any]) -> None: self.adam_spin_quadrupole_weight = np.asarray( variables["adam_spin_quadrupole_weight"], dtype=precision_dtype ) + self.spin_gate = np.asarray(variables["spin_gate"], dtype=precision_dtype) def set_spin_reference(self, reference: np.ndarray) -> None: """Store the per-type reference magnitudes. diff --git a/deepmd/kernels/cuda/dpa4c/graph_compress.py b/deepmd/kernels/cuda/dpa4c/graph_compress.py index a551b2beea..836cd71547 100644 --- a/deepmd/kernels/cuda/dpa4c/graph_compress.py +++ b/deepmd/kernels/cuda/dpa4c/graph_compress.py @@ -201,6 +201,26 @@ def has_spin(self) -> bool: """Return whether the compiled profile carries the spin families.""" return self.spin_channels > 0 + @property + def spin_slice(self) -> slice: + """Return the descriptor columns holding the spin invariants. + + The layout closes with the spin invariants, the two moment divisors + and the center type tail, so the block is addressed from the end and + is empty for a spin-free profile. + """ + stop = self.output_width - 2 - self.degree_channels[0] + vector_width = 1 + 2 * self.spin_channels + width = ( + 0 + if not self.has_spin + else vector_width * (vector_width + 1) // 2 + + 2 + + 2 * self.degree_channels[2] + + 2 * self.spin_channels + ) + return slice(stop - width, stop) + def descriptor_profile( channels: int, @@ -628,6 +648,17 @@ def build_compression_artifacts( output_inv_std = torch.reciprocal( descriptor.stddev.to(device=device, dtype=torch.float32) ) + if profile.has_spin: + # The operator assembles the spin invariants without reading the + # branch gate, which the portable path applies to the calibrated + # block. Scaling the inverse deviation of those columns carries + # the gate exactly, in the forward and in the backward alike: + # the kernel pulls every output cotangent back through that same + # array. A closed gate is a zero slope, so the gate never + # restricts what may be compressed. + output_inv_std[profile.spin_slice] *= float( + torch.as_tensor(descriptor.spin.spin_gate).reshape(()) + ) records, coupling_entry, coupling_value = coupling_records( profile.channels, diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py index c94158808e..c7d76b2d02 100644 --- a/deepmd/pt_expt/descriptor/dpa4c.py +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -41,6 +41,7 @@ "SpinChannels": ( "adam_spin_vector_weight", "adam_spin_quadrupole_weight", + "spin_gate", ), } diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index b053d0f913..e78be39250 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -54,6 +54,76 @@ def enable_hessian(self) -> None: self.requires_hessian("energy") self._hessian_enabled = True + def _to_public_keys( + self, + model_ret: dict[str, torch.Tensor], + do_atomic_virial: bool, + ) -> dict[str, torch.Tensor]: + """Rename one ``call_common`` result to the public energy-model keys. + + The renaming is a property of the output definition, not of the node + axis, so it serves the rectangular and the ragged entry alike. + + Parameters + ---------- + model_ret : dict[str, torch.Tensor] + A ``call_common`` result, in internal ``_`` keys. + do_atomic_virial : bool + Whether the per-atom virial was requested and should be carried. + + Returns + ------- + dict[str, torch.Tensor] + The same tensors under the public names. + """ + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy"): + model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy"): + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial: + model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) + for key in ("mask", "n_node"): + if key in model_ret: + model_predict[key] = model_ret[key] + if self.atomic_output_def()["energy"].r_hessian: + model_predict["hessian"] = model_ret["energy_derv_r_derv_r"].squeeze(-3) + return model_predict + + def _translate_eager_call( + self, + model_ret: dict[str, torch.Tensor], + atype: torch.Tensor, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + """Translate internal energy outputs at the public model boundary. + + Parameters + ---------- + model_ret : dict[str, torch.Tensor] + Result returned by a ``call_common`` entry. + atype : torch.Tensor + Atom types on the same node axis as the atomic outputs. + do_atomic_virial : bool, default: False + Whether the per-atom virial was requested. + + Returns + ------- + dict[str, torch.Tensor] + Public model outputs. + + Notes + ----- + Native-spin models override this translation to add ``force_mag`` and + ``mask_mag``. Keeping the dispatch here lets rectangular and ragged + forwards share one implementation without duplicating a spin-specific + ``forward_ragged``. + """ + del atype + return self._to_public_keys(model_ret, do_atomic_virial) + def forward_lower_canonical_graph( self, atype: torch.Tensor, @@ -297,45 +367,11 @@ def forward( do_atomic_virial=do_atomic_virial, neighbor_list=neighbor_list, ) - return self._to_public_keys(model_ret, do_atomic_virial) - - def _to_public_keys( - self, - model_ret: dict[str, torch.Tensor], - do_atomic_virial: bool, - ) -> dict[str, torch.Tensor]: - """Rename one ``call_common`` result to the public energy-model keys. - - The renaming is a property of the output definition, not of the node - axis, so it serves the rectangular and the ragged entry alike. - - Parameters - ---------- - model_ret : dict[str, torch.Tensor] - A ``call_common`` result, in internal ``_`` keys. - do_atomic_virial : bool - Whether the per-atom virial was requested and should be carried. - - Returns - ------- - dict[str, torch.Tensor] - The same tensors under the public names. - """ - model_predict = {} - model_predict["atom_energy"] = model_ret["energy"] - model_predict["energy"] = model_ret["energy_redu"] - if self.do_grad_r("energy"): - model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) - if self.do_grad_c("energy"): - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial: - model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) - for key in ("mask", "n_node"): - if key in model_ret: - model_predict[key] = model_ret[key] - if self.atomic_output_def()["energy"].r_hessian: - model_predict["hessian"] = model_ret["energy_derv_r_derv_r"].squeeze(-3) - return model_predict + return self._translate_eager_call( + model_ret, + atype, + do_atomic_virial=do_atomic_virial, + ) def forward_ragged( self, @@ -347,6 +383,7 @@ def forward_ragged( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Evaluate the energy model over a batch whose node axis is flat. @@ -373,6 +410,8 @@ def forward_ragged( Whether to return per-atom virials. charge_spin : torch.Tensor or None, optional Frame-level charge and spin conditioning with shape ``(nf, 2)``. + spin : torch.Tensor or None, optional + Per-atom native spin with shape ``(N, 3)``. Returns ------- @@ -389,8 +428,13 @@ def forward_ragged( aparam=aparam, do_atomic_virial=do_atomic_virial, charge_spin=charge_spin, + spin=spin, + ) + return self._translate_eager_call( + model_ret, + atype, + do_atomic_virial=do_atomic_virial, ) - return self._to_public_keys(model_ret, do_atomic_virial) def forward_lower( self, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 82193ba339..9dae115c03 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1142,92 +1142,6 @@ def _forward_eager( ) return self.original_model(coord, atype, **kwargs) - def forward_ragged( - self, - coord: torch.Tensor, - atype: torch.Tensor, - n_node: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - charge_spin: torch.Tensor | None = None, - spin: torch.Tensor | None = None, - ) -> dict[str, torch.Tensor]: - """Compiled forward over a batch whose node axis is already flat. - - The compiled lower works on that axis in either case -- its trace keeps - the frame count, the node count and the edge count as independent - symbols -- so a ragged batch simply skips the padding round trip the - rectangular :meth:`forward` performs around it. - - Parameters - ---------- - coord : torch.Tensor - Local coordinates with shape ``(N, 3)``, frame-major over ``n_node``. - atype : torch.Tensor - Local atom types with shape ``(N,)``. - n_node : torch.Tensor - Atoms per frame with shape ``(nf,)``. - box : torch.Tensor or None, optional - Simulation cell, ``(nf, 3, 3)`` or ``(nf, 9)``. - fparam : torch.Tensor or None, optional - Frame parameters with shape ``(nf, ndf)``. - aparam : torch.Tensor or None, optional - Atomic parameters with shape ``(N, nda)``. - do_atomic_virial : bool, default: False - Whether to return per-atom virials. - charge_spin : torch.Tensor or None, optional - Frame-level charge and spin conditioning with shape ``(nf, 2)``. - spin : torch.Tensor or None, optional - Native spin with shape ``(N, 3)``. - - Returns - ------- - dict[str, torch.Tensor] - Public model keys; per-atom entries keep the flat axis. - - Raises - ------ - NotImplementedError - If the model reads a rectangular node axis, which cannot represent - frames of unequal atom count without padding. - """ - if not self.training and not self._compile_eval: - return self._forward_eager( - coord, - atype, - box, - fparam, - aparam, - do_atomic_virial, - charge_spin, - spin, - n_node, - ) - del do_atomic_virial - if self._graph_eligible is None: - self._graph_eligible = model_uses_graph_lower(self.original_model) - if not self._graph_eligible: - raise NotImplementedError( - "a flat node axis requires a model whose descriptor reads one; " - "this model compiles the dense (nlist) lower, whose batches " - "must be padded to a common atom count" - ) - return self._forward_graph( - coord, - atype, - box, - fparam, - aparam, - charge_spin, - spin, - int(n_node.shape[0]), - 0, - self.original_model.get_rcut(), - n_node=n_node, - ) - def forward( self, coord: torch.Tensor, @@ -1474,6 +1388,93 @@ def forward( out[key] = val return out + def forward_ragged( + self, + coord: torch.Tensor, + atype: torch.Tensor, + n_node: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Compiled forward over a batch whose node axis is already flat. + + The compiled lower works on that axis in either case -- its trace keeps + the frame count, the node count and the edge count as independent + symbols -- so a ragged batch simply skips the padding round trip the + rectangular :meth:`forward` performs around it. + + Parameters + ---------- + coord : torch.Tensor + Local coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + atype : torch.Tensor + Local atom types with shape ``(N,)``. + n_node : torch.Tensor + Atoms per frame with shape ``(nf,)``. + box : torch.Tensor or None, optional + Simulation cell, ``(nf, 3, 3)`` or ``(nf, 9)``. + fparam : torch.Tensor or None, optional + Frame parameters with shape ``(nf, ndf)``. + aparam : torch.Tensor or None, optional + Atomic parameters with shape ``(N, nda)``. + do_atomic_virial : bool, default: False + Whether to return per-atom virials. + charge_spin : torch.Tensor or None, optional + Frame-level charge and spin conditioning with shape ``(nf, 2)``. + spin : torch.Tensor or None, optional + Native spin with shape ``(N, 3)``. Virtual-atom spin models do + not expose this ragged entry. + + Returns + ------- + dict[str, torch.Tensor] + Public model keys; per-atom entries keep the flat axis. + + Raises + ------ + NotImplementedError + If the model reads a rectangular node axis, which cannot represent + frames of unequal atom count without padding. + """ + if not self.training and not self._compile_eval: + return self._forward_eager( + coord, + atype, + box, + fparam, + aparam, + do_atomic_virial, + charge_spin, + spin, + n_node, + ) + del do_atomic_virial + if self._graph_eligible is None: + self._graph_eligible = model_uses_graph_lower(self.original_model) + if not self._graph_eligible: + raise NotImplementedError( + "a flat node axis requires a model whose descriptor reads one; " + "this model compiles the dense (nlist) lower, whose batches " + "must be padded to a common atom count" + ) + return self._forward_graph( + coord, + atype, + box, + fparam, + aparam, + charge_spin, + spin, + int(n_node.shape[0]), + 0, + self.original_model.get_rcut(), + n_node=n_node, + ) + def _forward_graph( self, coord: torch.Tensor, diff --git a/source/op/pt/dpa4c_graph_compress_kernel.cuh b/source/op/pt/dpa4c_graph_compress_kernel.cuh index b4a62b6c1d..c2f94f33f9 100644 --- a/source/op/pt/dpa4c_graph_compress_kernel.cuh +++ b/source/op/pt/dpa4c_graph_compress_kernel.cuh @@ -904,10 +904,11 @@ __device__ __forceinline__ void accumulate_coupling_gradient( // Four independent lane groups share one warp. An incomplete final block // aliases inactive groups to the last valid node so every thread reaches each // block-wide barrier; stores from those groups are suppressed. -template -__global__ __launch_bounds__(Profile::Threads, - 2) void node_backward_kernel(Arguments args) { - using P = Profile; +template +__global__ __launch_bounds__( + Profile::Threads, + 2) void node_backward_kernel(Arguments args) { + using P = Profile; constexpr int MaxComponents = 9; const int thread = threadIdx.x; const int group = thread / P::NodeWidth; @@ -1957,12 +1958,63 @@ template struct BackwardLauncher { + // Pick the node-group width from what the running device grants the two + // compiled kernels rather than from its architecture: the six shared arrays + // of the node backward scale with the group count and the scalar width, so + // the same source is shared-memory bound on a part with a 100 KB budget per + // multiprocessor and register bound on one with more than twice that. + // + // Widening the group halves that footprint and so buys resident warps, but + // every lane of a group reloads its node's shared state, so it also doubles + // the redundant load work. Measured across the scalar widths, the trade pays + // only where the narrow variant is starved of warps outright: at a quarter of + // the device's warp capacity the wide variant wins by 2.1 % of the whole + // backward, at a third the two are within noise, and above that the narrow + // variant wins by up to 3.6 %. The threshold is therefore stated against the + // device's own warp capacity, and the query is made once per configuration. + static bool prefer_wide_nodes() { + static const bool wide = [] { + using Narrow = Profile; + using Wide = Profile; + int narrow_blocks = 0; + int wide_blocks = 0; + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &narrow_blocks, + node_backward_kernel, + Narrow::Threads, 0); + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &wide_blocks, + node_backward_kernel, + Wide::Threads, 0); + int device = 0; + cudaGetDevice(&device); + int threads_per_sm = 0; + cudaDeviceGetAttribute(&threads_per_sm, + cudaDevAttrMaxThreadsPerMultiProcessor, device); + const int warps_per_sm = threads_per_sm / kWarpSize; + const int narrow_warps = narrow_blocks * (Narrow::Threads / kWarpSize); + return wide_blocks > narrow_blocks && narrow_warps * 4 < warps_per_sm; + }(); + return wide; + } + + template + static void launch_node_backward(const Arguments& args, cudaStream_t stream) { + using NodeProfile = Profile; + const int node_blocks = + static_cast((args.node_count + NodeProfile::NodeGroups - 1) / + NodeProfile::NodeGroups); + node_backward_kernel + <<>>(args); + } + static void run(const Arguments& args, cudaStream_t stream) { using P = Profile; - const int node_blocks = - static_cast((args.node_count + P::NodeGroups - 1) / P::NodeGroups); - node_backward_kernel - <<>>(args); + if (prefer_wide_nodes()) { + launch_node_backward(args, stream); + } else { + launch_node_backward(args, stream); + } edge_backward_kernel <<(args.node_count), P::Threads, 0, stream>>>(args); // The reserved edge slots beyond the physical count are only known on the diff --git a/source/op/pt/dpa4c_graph_compress_launch.h b/source/op/pt/dpa4c_graph_compress_launch.h index be885ac78f..70e52f1bd6 100644 --- a/source/op/pt/dpa4c_graph_compress_launch.h +++ b/source/op/pt/dpa4c_graph_compress_launch.h @@ -204,6 +204,19 @@ constexpr int coupling_record_count(int lmax) { // recompute that a wider group amortizes. A group wider than the scalar width // leaves a tile empty and does not compile, which bounds the narrow profiles // from above. +// Node-group width candidates. The node kernels hold six shared arrays whose +// extent is the product of the resident node groups and the scalar width: the +// three geometric quantities of the recompute and their three cotangents. That +// footprint, not the register file, is what bounds the backward's resident +// blocks once the scalar width grows, and how soon it binds depends on the +// device's shared-memory capacity rather than on its architecture generation. +// Both widths are compiled and :func:`BackwardLauncher::prefer_wide_nodes` +// picks between them from what the occupancy API reports for the running +// device. A width is a lane count, so it must be a power of two and must not +// exceed the warp. +constexpr int kNodeLanesNarrow = 8; +constexpr int kNodeLanesWide = 16; + template struct EdgeMap; @@ -246,7 +259,10 @@ struct EdgeMap<128> { // ``HasSpin`` is deliberately without a default. It changes the moment // layout and the descriptor width, so an instantiation that omits it would // silently read a spin-free layout out of a spin-conditioned buffer. -template +// ``NodeLanes`` overrides the node-group width; zero selects the narrow +// candidate. Only the node kernels instantiate the override, so the edge +// kernels and every existing use of ``Profile`` are unaffected. +template struct Profile { static constexpr int C0 = Channels; static constexpr int C1 = degree_one_width(Channels); @@ -344,7 +360,8 @@ struct Profile { HasSpin ? EdgeMap::SpinForward : EdgeMap::Forward; static constexpr int BackwardEdgeWidth = HasSpin ? EdgeMap::SpinBackward : EdgeMap::Backward; - static constexpr int NodeWidth = 8; + static constexpr int NodeWidth = + NodeLanes != 0 ? NodeLanes : kNodeLanesNarrow; static constexpr int NodeGroups = kWarpSize / NodeWidth; static constexpr int Threads = kWarpSize; diff --git a/source/tests/common/dpmodel/test_descriptor_dpa4c.py b/source/tests/common/dpmodel/test_descriptor_dpa4c.py index 4c7c583cfd..9832dabd8c 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa4c.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa4c.py @@ -21,6 +21,9 @@ enumerate_degree_triples, packed_l2_to_stf, ) +from deepmd.dpmodel.descriptor.dpa4c_nn.spin import ( + SpinChannels, +) from deepmd.dpmodel.utils import ( neighbor_graph, ) @@ -778,7 +781,12 @@ def test_automatic_profiles_and_output_dimensions( def make_spin_descriptor(**overrides: Any) -> DescrptDPA4C: - """Build a descriptor whose first atom type carries a magnetic moment.""" + """Build a descriptor whose first atom type carries a magnetic moment. + + The branch gate is opened, since a fresh descriptor starts spin-free by + design and the tests below are about the branch behind the gate. The gate + itself is covered by :class:`TestDPA4CSpinGate`. + """ config: dict[str, Any] = { "rcut": 3.0, "ntypes": 2, @@ -790,7 +798,10 @@ def make_spin_descriptor(**overrides: Any) -> DescrptDPA4C: "use_spin": [True, False], } config.update(overrides) - return DescrptDPA4C(**config) + descriptor = DescrptDPA4C(**config) + if descriptor.spin is not None: + descriptor.spin.spin_gate[...] = 1.0 + return descriptor def spin_reference_terms( @@ -882,6 +893,89 @@ def vector_gram_selectors(descriptor: DescrptDPA4C) -> dict[str, np.ndarray]: } +class TestDPA4CSpinGate: + """The scalar gate that carries the whole spin branch. + + Activating a magnetic type on a spin-free pretraining releases weights + that never received a gradient, so the branch must start at the zero + function instead. No weight inside the branch can do that: the families + reach the fitting network by several routes and at two spin orders, and a + factor on the conditioned moment would enter the invariants quadratically + and leave zero a stationary point. + """ + + def setup_method(self) -> None: + self.descriptor = make_spin_descriptor() + rng = np.random.default_rng(5) + self.spin = rng.normal(size=(SPIN_ATYPE.size, 3)) + self.graph = neighbor_graph.build_neighbor_graph( + SPIN_COORD, + SPIN_ATYPE, + None, + self.descriptor.get_rcut(), + ) + self.atype = SPIN_ATYPE.reshape(-1) + + def evaluate(self, spin: np.ndarray | None) -> np.ndarray: + return self.descriptor.call_graph(self.graph, self.atype, spin=spin)[0] + + def test_fresh_descriptor_starts_spin_free(self) -> None: + """A freshly built descriptor closes the gate.""" + fresh = DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=8, + lmax=2, + n_radial=4, + precision="float64", + seed=23, + use_spin=[True, False], + ) + np.testing.assert_array_equal(fresh.spin.spin_gate, 0.0) + + def spin_block(self, output: np.ndarray) -> np.ndarray: + start = self.descriptor.readout.get_dim_out() + return output[:, start : start + self.descriptor.spin.get_dim_out()] + + def test_closed_gate_erases_the_whole_branch(self) -> None: + """A closed gate zeroes every family, for arbitrary moments. + + Zero moments are not the spin-free reference here: the magnetic + effective coordination weighs the per-type mask rather than the + moment, so it survives ``s = 0`` and naming a type magnetic moves the + descriptor on its own. Only the gate removes that term as well. + """ + assert np.any(self.spin_block(self.evaluate(np.zeros_like(self.spin)))) + self.descriptor.spin.spin_gate[...] = 0.0 + for moments in (self.spin, np.zeros_like(self.spin)): + np.testing.assert_array_equal(self.spin_block(self.evaluate(moments)), 0.0) + + def test_invariants_are_linear_in_the_gate(self) -> None: + """Linearity is what leaves the closed gate a nonzero gradient. + + A factor applied to the conditioned moment instead would reach the + degree-one Grams squared and the quadrupole Grams to the fourth + power, so its derivative would vanish with the gate itself. + """ + unit = self.spin_block(self.evaluate(self.spin)) + for gate in (0.25, -1.5, 3.0): + self.descriptor.spin.spin_gate[...] = gate + np.testing.assert_allclose( + self.spin_block(self.evaluate(self.spin)), + gate * unit, + rtol=1e-12, + atol=1e-14, + ) + + def test_serialization_carries_the_gate(self) -> None: + """The gate is trained state and therefore round-trips.""" + self.descriptor.spin.spin_gate[...] = 0.42 + payload = self.descriptor.spin.serialize() + np.testing.assert_allclose( + SpinChannels.deserialize(payload).spin_gate, 0.42, rtol=1e-12 + ) + + class TestDPA4CSpin: def setup_method(self) -> None: self.descriptor = make_spin_descriptor() diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py index 9e3f41a24a..f52b4f346a 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -428,6 +428,63 @@ def test_coincident_edge_has_finite_third_derivative(self) -> None: ) +class TestDPA4CSpinGate: + """Torch-side contracts of the spin branch gate.""" + + def make_descriptor(self) -> DescrptDPA4C: + return DescrptDPA4C( + rcut=3.0, + ntypes=2, + channels=8, + lmax=2, + n_radial=4, + precision="float64", + seed=23, + use_spin=[True, False], + ).to(env.DEVICE) + + def test_fresh_descriptor_starts_spin_free(self) -> None: + assert float(self.make_descriptor().spin.spin_gate.detach()) == 0.0 + + def test_closed_gate_still_receives_a_gradient(self) -> None: + """Zero is a starting point, not a fixed point. + + The invariants are linear in the gate, so its gradient there is the + branch itself. A factor on the conditioned moment would reach the + Grams at second and fourth order and could never reopen. + """ + descriptor = self.make_descriptor() + generator = torch.Generator(device="cpu").manual_seed(7) + coord = ( + torch.randn(1, 6, 3, dtype=torch.float64, generator=generator).to( + env.DEVICE + ) + * 1.4 + ) + atype = torch.tensor([[0, 1, 0, 1, 0, 1]], dtype=torch.long, device=env.DEVICE) + spin = torch.randn(6, 3, dtype=torch.float64, generator=generator).to( + env.DEVICE + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + graph = build_neighbor_graph(coord, atype, None, 3.0) + output, _ = descriptor.call_graph(graph, atype.reshape(-1), spin=spin) + output.sum().backward() + gradient = descriptor.spin.spin_gate.grad + assert gradient is not None + assert float(gradient.abs().max()) > 1e-8 + + def test_a_stored_gate_round_trips(self) -> None: + descriptor = self.make_descriptor() + with torch.no_grad(): + descriptor.spin.spin_gate.fill_(0.37) + restored = self.make_descriptor() + restored.load_state_dict(descriptor.state_dict()) + assert float(restored.spin.spin_gate.detach()) == pytest.approx(0.37) + + class TestDPA4CSpin: """Torch-side contracts of the native spin branch.""" @@ -443,6 +500,10 @@ def setup_method(self) -> None: use_spin=[True, False], ).to(env.DEVICE) self.descriptor.eval() + # A fresh descriptor starts spin-free by design; these tests are about + # the branch behind the gate, which ``TestDPA4CSpinGate`` covers. + with torch.no_grad(): + self.descriptor.spin.spin_gate.fill_(1.0) generator = torch.Generator(device="cpu").manual_seed(5) self.coord = ( torch.randn(1, 6, 3, dtype=torch.float64, generator=generator).to( @@ -475,7 +536,9 @@ def test_spin_arrays_are_optimizer_visible(self) -> None: names = {name for name, _ in self.descriptor.named_parameters()} assert "spin.adam_spin_vector_weight" in names assert "spin.adam_spin_quadrupole_weight" in names - # The gate and the reference are state, not learned quantities. + assert "spin.spin_gate" in names + # The per-type mask and the reference magnitudes are state, not + # learned quantities. buffers = {name for name, _ in self.descriptor.named_buffers()} assert {"spin.spin_mask", "spin.spin_reference"} <= buffers diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py index 2ab071467d..c3bb5e0476 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py @@ -1216,7 +1216,9 @@ def _build_spin_descriptor( A reference magnitude other than one makes the conditioning factor visible in the magnetic gradient, so a missing chain factor shows up as a constant - ratio rather than cancelling. + ratio rather than cancelling. The branch gate is opened, since a fresh + descriptor starts spin-free by design and these tests are about the branch + behind the gate; the gate itself is covered in ``test_dpa4c.py``. """ descriptor = ( DescrptDPA4C( @@ -1234,6 +1236,8 @@ def _build_spin_descriptor( .eval() ) descriptor.spin.set_spin_reference(np.array([1.7, 1.0, 1.0])) + with torch.no_grad(): + descriptor.spin.spin_gate.fill_(1.0) return descriptor diff --git a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py index 7ed04e249d..279b52cd7a 100644 --- a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py @@ -4,6 +4,7 @@ SimpleNamespace, ) +import numpy as np import pytest import torch @@ -119,6 +120,227 @@ def test_graph_lower_energy_force_are_finite() -> None: assert torch.isfinite(result[key]).all(), key +def test_phantom_atoms_leave_energy_and_force_unchanged() -> None: + """A mixed-nloc batch's padding must not perturb its real atoms. + + Frames of unequal atom count only share a batch when the shorter ones are + padded to a rectangular shape, with the padded slots marked ``atype = -1``. + Such a phantom atom stands for no physical site: the graph builders drop it + from the edge set and the atomic model zeroes its output, so the padded + batch has to reproduce, frame by frame, what the frames give on their own. + """ + torch.manual_seed(1234) + model = get_model(_config()).to(env.DEVICE) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + model.eval() + + rng = np.random.default_rng(0) + nlocs = (4, 7, 3) + pad_nloc = max(nlocs) + box = (np.eye(3) * 10.0).reshape(9) + coords = [rng.uniform(0.0, 6.0, (nloc, 3)) for nloc in nlocs] + atypes = [rng.integers(0, 2, nloc) for nloc in nlocs] + + padded_coord = np.zeros((len(nlocs), pad_nloc, 3)) + padded_atype = np.full((len(nlocs), pad_nloc), -1, dtype=np.int64) + for index, (frame_coord, frame_atype) in enumerate( + zip(coords, atypes, strict=True) + ): + padded_coord[index, : len(frame_atype)] = frame_coord + padded_atype[index, : len(frame_atype)] = frame_atype + + def run(coord: np.ndarray, atype: np.ndarray) -> dict[str, torch.Tensor]: + nframes = coord.shape[0] + return model( + torch.tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(atype, dtype=torch.long, device=env.DEVICE), + box=torch.tensor( + np.tile(box, (nframes, 1)), dtype=torch.float64, device=env.DEVICE + ), + ) + + batched = run(padded_coord, padded_atype) + for index, (frame_coord, frame_atype) in enumerate( + zip(coords, atypes, strict=True) + ): + alone = run(frame_coord[None], frame_atype[None].astype(np.int64)) + torch.testing.assert_close( + batched["energy"].reshape(-1)[index], + alone["energy"].reshape(-1)[0], + atol=1.0e-12, + rtol=1.0e-12, + ) + torch.testing.assert_close( + batched["force"][index, : len(frame_atype)], + alone["force"][0], + atol=1.0e-12, + rtol=1.0e-12, + ) + # Padded slots receive no gradient at all. + phantom = torch.tensor(padded_atype, device=env.DEVICE) < 0 + assert bool(phantom.any()), "fixture must exercise padding" + assert bool(torch.all(batched["force"][phantom] == 0.0)) + + +def test_padding_never_reaches_the_network() -> None: + """The node axis the graph lower receives holds the real atoms alone. + + The equivalence test above holds whether or not the phantom atoms are + dropped, since the atomic model masks their output either way. What is + asserted here is that they are dropped: the padded slots must cost no + descriptor, no fitting-net evaluation and no gradient, which is the whole + point of packing frames of unequal atom count into one batch. + """ + model = get_model(_config()).to(env.DEVICE).eval() + + seen: list[int] = [] + lower = model.forward_common_lower_graph + + def spy(atype, *args, **kwargs): + seen.append(int(atype.shape[0])) + return lower(atype, *args, **kwargs) + + model.forward_common_lower_graph = spy + + rng = np.random.default_rng(5) + nlocs = (4, 7, 3) + pad_nloc = max(nlocs) + padded_coord = np.zeros((len(nlocs), pad_nloc, 3)) + padded_atype = np.full((len(nlocs), pad_nloc), -1, dtype=np.int64) + for index, nloc in enumerate(nlocs): + padded_coord[index, :nloc] = rng.uniform(0.0, 6.0, (nloc, 3)) + padded_atype[index, :nloc] = rng.integers(0, 2, nloc) + + out = model( + torch.tensor(padded_coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(padded_atype, dtype=torch.long, device=env.DEVICE), + box=torch.tensor( + np.tile((np.eye(3) * 10.0).reshape(9), (len(nlocs), 1)), + dtype=torch.float64, + device=env.DEVICE, + ), + ) + + assert seen == [sum(nlocs)], ( + f"the lower saw {seen} nodes; the batch holds {sum(nlocs)} real atoms " + f"padded to {len(nlocs) * pad_nloc} slots" + ) + # The public output still carries the padded shape the callers expect. + assert out["force"].shape == (len(nlocs), pad_nloc, 3) + + +def test_compiled_lower_accepts_a_compacted_node_axis() -> None: + """The compiled artifact must not carry ``N == nframes * nloc`` as a guard. + + Its trace is taken on a uniform system, where the flat node axis happens to + be the product of the frame count and the atom count. Dropping the padding + breaks that relation, so this exercises the compiled lower on a batch where + it no longer holds, and holds the result against the eager graph path, + which takes the same compaction. + """ + from deepmd.pt_expt.train.training import ( + _CompiledModel, + _get_model_structure_key, + ) + + torch.manual_seed(0) + config = _config() + config["descriptor"]["channels"] = 8 + config["fitting_net"]["neuron"] = [8, 8] + model = get_model(config).to(env.DEVICE).train() + compiled = _CompiledModel(model, _get_model_structure_key(model)) + + rng = np.random.default_rng(0) + nlocs = (4, 7, 3) + pad_nloc = max(nlocs) + padded_coord = np.zeros((len(nlocs), pad_nloc, 3)) + padded_atype = np.full((len(nlocs), pad_nloc), -1, dtype=np.int64) + for index, nloc in enumerate(nlocs): + padded_coord[index, :nloc] = rng.uniform(0.0, 6.0, (nloc, 3)) + padded_atype[index, :nloc] = rng.integers(0, 2, nloc) + + args = ( + torch.tensor(padded_coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(padded_atype, dtype=torch.long, device=env.DEVICE), + torch.tensor( + np.tile((np.eye(3) * 10.0).reshape(1, 3, 3), (len(nlocs), 1, 1)), + dtype=torch.float64, + device=env.DEVICE, + ), + ) + got = compiled(*args) + expected = model(*args) + + assert got["force"].shape == (len(nlocs), pad_nloc, 3) + torch.testing.assert_close(got["energy"], expected["energy"]) + torch.testing.assert_close(got["force"], expected["force"]) + phantom = args[1] < 0 + assert bool(torch.all(got["force"][phantom] == 0.0)) + + +def test_ragged_and_padded_batches_agree() -> None: + """The two layouts are two spellings of one batch, so they must agree. + + This is the invariant the whole flat-node-axis path rests on: concatenating + the frames rather than padding them to a common width changes how the batch + is stored, and nothing about the physics it describes. + """ + torch.manual_seed(11) + model = get_model(_config()).to(env.DEVICE) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.1) + model.eval() + + rng = np.random.default_rng(0) + nlocs = (4, 7, 3) + pad_nloc, boxlen = max(nlocs), 10.0 + padded_coord = np.zeros((len(nlocs), pad_nloc, 3)) + padded_atype = np.full((len(nlocs), pad_nloc), -1, dtype=np.int64) + flat_coord, flat_atype = [], [] + for index, nloc in enumerate(nlocs): + coord = rng.uniform(0.0, 6.0, (nloc, 3)) + atype = rng.integers(0, 2, nloc) + padded_coord[index, :nloc] = coord + padded_atype[index, :nloc] = atype + flat_coord.append(coord) + flat_atype.append(atype) + + padded = model( + torch.tensor(padded_coord, dtype=torch.float64, device=env.DEVICE), + torch.tensor(padded_atype, dtype=torch.long, device=env.DEVICE), + box=torch.tensor( + np.tile((np.eye(3) * boxlen).reshape(9), (len(nlocs), 1)), + dtype=torch.float64, + device=env.DEVICE, + ), + ) + ragged = model.forward_ragged( + torch.tensor( + np.concatenate(flat_coord), dtype=torch.float64, device=env.DEVICE + ), + torch.tensor(np.concatenate(flat_atype), dtype=torch.long, device=env.DEVICE), + torch.tensor(nlocs, dtype=torch.long, device=env.DEVICE), + torch.tensor( + np.tile(np.eye(3)[None] * boxlen, (len(nlocs), 1, 1)), + dtype=torch.float64, + device=env.DEVICE, + ), + ) + + torch.testing.assert_close(ragged["energy"], padded["energy"]) + offset = 0 + for index, nloc in enumerate(nlocs): + torch.testing.assert_close( + ragged["force"][offset : offset + nloc], + padded["force"][index, :nloc], + ) + offset += nloc + assert offset == ragged["force"].shape[0], "the ragged axis holds real atoms only" + + def test_graph_force_loss_trains_descriptor() -> None: model = get_model(_config()).to(env.DEVICE).train() result = _run_graph(model) @@ -366,8 +588,10 @@ def _spin_sample(model: torch.nn.Module) -> tuple: @pytest.mark.skipif( not torch.cuda.is_available(), reason="the fused spin path is CUDA only" ) +@pytest.mark.parametrize("gate", [0.8, 0.0]) def test_compressed_spin_lowers_match_autograd( monkeypatch: pytest.MonkeyPatch, + gate: float, ) -> None: """Both fused lowers reproduce the autograd magnetic force. @@ -376,12 +600,21 @@ def test_compressed_spin_lowers_match_autograd( neighbour half is emitted per edge and reduced onto source nodes. Checking it against the autograd lower covers the fused generic composition and the compact canonical deployment path in the same comparison. + + The operator assembles the spin invariants without reading the branch + gate, which the portable path applies to the calibrated block; compression + carries it as a factor on the inverse deviation of those columns. A + non-unit value is what makes this comparison check that fold, and the + closed gate pins the case an affine mean-and-deviation fold could not + express -- a trained gate may legitimately reach zero. """ import numpy as np model = get_model(_spin_config()).to(env.DEVICE).eval() descriptor = model.get_descriptor() descriptor.spin.set_spin_reference(np.array([1.7, 1.0, 1.0])) + with torch.no_grad(): + descriptor.spin.spin_gate.fill_(gate) graph, atype, spin = _spin_sample(model) descriptor.enable_compression(min_nbor_dist=0.5) diff --git a/source/tests/pt_expt/test_finetune.py b/source/tests/pt_expt/test_finetune.py index bdba008ec8..7413743481 100644 --- a/source/tests/pt_expt/test_finetune.py +++ b/source/tests/pt_expt/test_finetune.py @@ -1023,6 +1023,18 @@ def test_finetune_from_pt2_use_pretrain_script(self) -> None: }, "dpa4_ener", ), + "dpa4c": ( + { + "type": "dpa4c", + "rcut": 3.0, + "channels": 8, + "lmax": 2, + "n_radial": 8, + "precision": "float64", + "seed": 17, + }, + "ener", + ), } From dab3a6e8a0783836f11cd6d75489a62e8d1215b8 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 14 Aug 2026 17:37:31 +0800 Subject: [PATCH 05/10] fix(argcheck): use backend support registry for DPA4C --- deepmd/utils/argcheck.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 14b768320a..84c0726a20 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -455,7 +455,7 @@ def descrpt_se_a_args() -> list[Argument]: @descrpt_args_plugin.register( "dpa4c", alias=["DPA4C"], - doc=doc_only_pt_expt_supported + doc=supported_backends("pt_expt") + "DPA4C is the compact and compressible degree-wise descriptor of the DPA4 family.", ) def descrpt_dpa4c_args() -> list[Argument]: From b0a7c0347b6ada130d5d1bec691aeee83111c842 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 09:41:33 +0000 Subject: [PATCH 06/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt_expt/train/validation.py | 4 +--- examples/water/dpa4c/input.json | 21 +++++++++++++++++---- source/op/pt/CMakeLists.txt | 4 ++-- 3 files changed, 20 insertions(+), 9 deletions(-) diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py index 1af9e454bc..3c13009400 100644 --- a/deepmd/pt_expt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -524,9 +524,7 @@ def _evaluate_system( test_data = data_system.get_test() natoms = int(test_data["type"].shape[1]) nframes = int(test_data["coord"].shape[0]) - include_virial = ( - data_system.pbc and bool(test_data.get("find_virial", 0.0)) - ) + include_virial = data_system.pbc and bool(test_data.get("find_virial", 0.0)) spin = ( test_data["spin"].reshape(nframes, -1) if self.profile.needs_spin else None ) diff --git a/examples/water/dpa4c/input.json b/examples/water/dpa4c/input.json index 3b5e5e80df..b658821082 100644 --- a/examples/water/dpa4c/input.json +++ b/examples/water/dpa4c/input.json @@ -1,7 +1,10 @@ { "_comment": "DPA4C energy-training example for the water dataset.", "model": { - "type_map": ["O", "H"], + "type_map": [ + "O", + "H" + ], "descriptor": { "type": "dpa4c", "_comment": "The Neo grade: channels and lmax fix every derived width, while radial_modes spends per-edge work without widening the per-atom state.", @@ -18,7 +21,11 @@ }, "fitting_net": { "_comment": "The fitting width paired with channels 32 by the Neo grade.", - "neuron": [192, 192, 192], + "neuron": [ + 192, + 192, + 192 + ], "resnet_dt": false, "activation_function": "silu", "precision": "float32", @@ -53,11 +60,17 @@ "stat_file": "./dpa4c.hdf5", "stat_file_mode": "update", "training_data": { - "systems": ["../data/data_0", "../data/data_1", "../data/data_2"], + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], "batch_size": 1 }, "validation_data": { - "systems": ["../data/data_3"], + "systems": [ + "../data/data_3" + ], "batch_size": 1, "numb_btch": 1 }, diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index adb3bd4aad..b1d6c31cc3 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -74,8 +74,8 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) target_compile_definitions(deepmd_op_pt PRIVATE DEEPMD_ENABLE_DPA1_HIGH_LMAX=1) endif() - # The compressed DPA1 and DPA4C kernels are instantiated one translation - # unit per channel width so their angular-degree and topology specializations + # The compressed DPA1 and DPA4C kernels are instantiated one translation unit + # per channel width so their angular-degree and topology specializations # compile in parallel. set_source_files_properties( ${DPA1_GRAPH_COMPRESS_KERNEL_SRC} dpa4c_graph_compress_c8.cu From 279b7405ee360644db49b9f8fbae4d308c13ba1c Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 14 Aug 2026 21:27:37 +0800 Subject: [PATCH 07/10] fix(dpa4c): harden compressed inference contracts --- deepmd/dpmodel/descriptor/dpa4c.py | 3 +- deepmd/dpmodel/loss/ener.py | 47 ++++----- deepmd/kernels/cuda/dpa1/canonical.py | 2 +- deepmd/kernels/cuda/dpa1/graph_compress.py | 2 +- .../kernels/cuda/dpa1/graph_energy_force.py | 2 +- deepmd/kernels/cuda/dpa4c/canonical.py | 17 +++- deepmd/kernels/cuda/dpa4c/graph_compress.py | 27 +++--- deepmd/kernels/cuda/edge_force_virial.py | 36 ++++--- deepmd/kernels/cuda/graph_fitting.py | 12 ++- deepmd/pt/loss/ener.py | 41 ++++---- deepmd/pt_expt/descriptor/dpa1.py | 17 ++-- deepmd/pt_expt/model/edge_transform_output.py | 2 +- deepmd/pt_expt/model/ener_model.py | 9 +- deepmd/pt_expt/model/make_model.py | 2 +- deepmd/pt_expt/utils/serialization.py | 17 +++- deepmd/utils/eval_metrics.py | 6 +- source/api_c/include/c_api.h | 8 +- source/api_c/include/c_api_internal.h | 11 +++ source/api_c/src/c_api.cc | 35 +++++-- source/api_c/tests/test_deepmd_exception.cc | 6 ++ source/api_cc/include/NativeSpinPTExpt.h | 28 +++--- source/api_cc/include/commonPT.h | 27 +++++- source/api_cc/src/DeepPotPTExpt.cc | 19 ++-- source/api_cc/src/DeepSpinPTExpt.cc | 14 +-- source/api_cc/src/NativeSpinPTExpt.cc | 49 +++++++--- source/api_cc/src/commonPTExpt.h | 14 ++- .../api_cc/tests/test_neighbor_list_data.cc | 11 ++- source/lmp/compact_canonical_graph_kokkos.h | 2 +- source/lmp/pair_deepmd_kokkos.cpp | 3 +- source/lmp/pair_dpa4spin.cpp | 14 ++- source/op/pt/dpa4c_graph_compress.cu | 29 ++++-- source/op/pt/edge_force_virial.cu | 63 ++++++------ source/op/pt/graph_fitting.cu | 53 ++++++++--- source/op/pt/graph_ops.h | 8 +- .../tests/common/dpmodel/test_loss_padding.py | 41 +++++++- source/tests/pt/test_loss_padding.py | 35 ++++++- .../pt_expt/descriptor/test_dpa1_cuda.py | 41 ++++++-- .../pt_expt/descriptor/test_dpa4c_cuda.py | 95 +++++++++++++++++-- .../pt_expt/model/test_dpa4c_graph_lower.py | 44 +++++++++ 39 files changed, 640 insertions(+), 252 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index d714a5f515..389dbd27f8 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -573,8 +573,7 @@ def call( This method exists for the common descriptor ABI and numerical reference tests. Production DPA4C execution uses :meth:`call_graph` - with a carry-all graph. A rectangular list at the internal compatibility - capacity is rejected because its completeness cannot be established. + with a carry-all graph. Parameters ---------- diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 1abe910a2f..de7b502c84 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -33,6 +33,15 @@ ) +def _huber_from_residual(residual: Array, delta: float = 1.0) -> Array: + xp = array_api_compat.array_namespace(residual) + abs_error = xp.abs(residual) + quadratic_loss = 0.5 * residual**2 + linear_loss = delta * (abs_error - 0.5 * delta) + loss = xp.where(abs_error <= delta, quadratic_loss, linear_loss) + return xp.mean(loss) + + def custom_huber_loss(predictions: Array, targets: Array, delta: float = 1.0) -> Array: r"""Return the mean Huber loss. @@ -45,13 +54,7 @@ def custom_huber_loss(predictions: Array, targets: Array, delta: float = 1.0) -> \delta(|e|-\tfrac12\delta),& |e|>\delta. \end{cases} """ - xp = array_api_compat.array_namespace(predictions, targets) - error = targets - predictions - abs_error = xp.abs(error) - quadratic_loss = 0.5 * error**2 - linear_loss = delta * (abs_error - 0.5 * delta) - loss = xp.where(abs_error <= delta, quadratic_loss, linear_loss) - return xp.mean(loss) + return _huber_from_residual(targets - predictions, delta) class EnergyLoss(Loss): @@ -498,13 +501,9 @@ def call( ) # [nf, nloc, 3] huber_ncomp = 3 else: - diff_3 = xp.reshape( - force_hat_reshape - force_reshape, - (*_node_shape, 3), - ) norm_2d = xp.reshape( xp.linalg.vector_norm( - xp.reshape(diff_3, (-1, 3)), axis=1 + xp.reshape(diff_f_3d, (-1, 3)), axis=1 ), _node_shape, ) @@ -532,21 +531,17 @@ def call( loss += pref_f * l2_force_loss else: if not self.f_use_norm: - l_huber_loss = custom_huber_loss( - xp.reshape(force, (-1,)), - xp.reshape(force_hat, (-1,)), + l_huber_loss = _huber_from_residual( + diff_f, delta=self._huber_delta_force, ) else: - force_diff_3 = xp.reshape( - force_hat_reshape - force_reshape, (-1, 3) - ) + force_diff_3 = xp.reshape(diff_f, (-1, 3)) force_diff_norm = xp.reshape( xp.linalg.vector_norm(force_diff_3, axis=1), (-1, 1) ) - l_huber_loss = custom_huber_loss( + l_huber_loss = _huber_from_residual( force_diff_norm, - xp.zeros_like(force_diff_norm), delta=self._huber_delta_force, ) loss += pref_f * l_huber_loss @@ -559,12 +554,10 @@ def call( if not self.f_use_norm: l1_force_masked = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3) else: - diff_3 = xp.reshape( - force_hat_reshape - force_reshape, - (*_node_shape, 3), - ) norm_2d = xp.reshape( - xp.linalg.vector_norm(xp.reshape(diff_3, (-1, 3)), axis=1), + xp.linalg.vector_norm( + xp.reshape(diff_f_3d, (-1, 3)), axis=1 + ), _node_shape, ) # One L2 norm per atom, hence one label per atom. @@ -579,9 +572,7 @@ def call( if not self.f_use_norm: l1_force_loss = xp.mean(xp.abs(diff_f)) else: - force_diff_3 = xp.reshape( - force_hat_reshape - force_reshape, (-1, 3) - ) + force_diff_3 = xp.reshape(diff_f, (-1, 3)) l1_force_loss = xp.mean( xp.linalg.vector_norm(force_diff_3, axis=1) ) diff --git a/deepmd/kernels/cuda/dpa1/canonical.py b/deepmd/kernels/cuda/dpa1/canonical.py index 54c0c8b2ea..997fd14215 100644 --- a/deepmd/kernels/cuda/dpa1/canonical.py +++ b/deepmd/kernels/cuda/dpa1/canonical.py @@ -443,7 +443,7 @@ def dpa1_canonical_compress_energy_force( graph.source_row_ptr, graph.source_order, graph.n_node, - graph.edge_vec.new_zeros(0, 3), + graph.edge_vec.new_empty(0), atype.shape[0], do_atomic_virial, ) diff --git a/deepmd/kernels/cuda/dpa1/graph_compress.py b/deepmd/kernels/cuda/dpa1/graph_compress.py index 8a62384ebe..fd9e734464 100644 --- a/deepmd/kernels/cuda/dpa1/graph_compress.py +++ b/deepmd/kernels/cuda/dpa1/graph_compress.py @@ -1022,7 +1022,7 @@ def dpa1_graph_compress_energy_force( graph.source_order, graph.source_row_ptr, graph.n_node, - edge_vec.new_zeros(0, 3), + edge_vec.new_empty(0), node_capacity, do_atomic_virial, ) diff --git a/deepmd/kernels/cuda/dpa1/graph_energy_force.py b/deepmd/kernels/cuda/dpa1/graph_energy_force.py index d3e627053c..40b5e536f4 100644 --- a/deepmd/kernels/cuda/dpa1/graph_energy_force.py +++ b/deepmd/kernels/cuda/dpa1/graph_energy_force.py @@ -276,7 +276,7 @@ def _cpu( source_order, source_row_ptr, n_node, - edge_vec.to(fprec).new_zeros(0, 3), + edge_vec.to(fprec).new_empty(0), node_capacity, do_atomic_virial, ) diff --git a/deepmd/kernels/cuda/dpa4c/canonical.py b/deepmd/kernels/cuda/dpa4c/canonical.py index 25504227d5..ce17c55983 100644 --- a/deepmd/kernels/cuda/dpa4c/canonical.py +++ b/deepmd/kernels/cuda/dpa4c/canonical.py @@ -64,9 +64,14 @@ def op_available() -> bool: "dpa4c_canonical_compress_backward_inplace", None, ) + energy_gradient = getattr( + torch.ops.deepmd, + "dpa4c_canonical_compress_energy_gradient", + None, + ) return all( isinstance(operator, torch._ops.OpOverloadPacket) - for operator in (forward, backward, backward_inplace) + for operator in (forward, backward, backward_inplace, energy_gradient) ) @@ -120,7 +125,7 @@ def _forward_fake( ) profile = descriptor_profile( - int(type_embedding.shape[1]), int(lmax), spin.numel() != 0 + int(type_embedding.shape[1]), int(lmax), spin.ndim == 2 ) nodes = atype.shape[0] descriptor = edge_vec.new_empty(nodes, profile.output_width, dtype=torch.float32) @@ -157,7 +162,7 @@ def _backward_shapes( Each absent output is allocated separately, because the schema declares three unannotated results and two of them may not share storage. """ - if spin.numel() == 0: + if spin.ndim != 2: return ( torch.empty_like(edge_vec), edge_vec.new_empty((0,), dtype=torch.float32), @@ -179,7 +184,7 @@ def _energy_gradient_fake( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: del source, destination_row_ptr spin = args[10] - has_spin = spin.numel() != 0 + has_spin = spin.ndim == 2 # Each absent output is allocated separately, because the schema declares # four unannotated results and no two of them may share storage. return ( @@ -194,7 +199,9 @@ def _energy_gradient_fake( ) -def _cpu_energy_gradient(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: +def _cpu_energy_gradient( + *args: Any, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Reference sequence of the fused operator, evaluated in one run.""" from deepmd.kernels.cuda.graph_fitting import _cpu_backward as fitting_backward from deepmd.kernels.cuda.graph_fitting import _cpu_forward as fitting_forward diff --git a/deepmd/kernels/cuda/dpa4c/graph_compress.py b/deepmd/kernels/cuda/dpa4c/graph_compress.py index 836cd71547..d1e460d756 100644 --- a/deepmd/kernels/cuda/dpa4c/graph_compress.py +++ b/deepmd/kernels/cuda/dpa4c/graph_compress.py @@ -179,6 +179,8 @@ class DescriptorProfile: Descriptor coordinate of the first degree-one Gram entry. bispectrum_base Descriptor coordinate of the first bispectrum entry. + spin_width + Width of the trailing native-spin invariant block. """ channels: int @@ -190,6 +192,7 @@ class DescriptorProfile: gram_base: int bispectrum_base: int spin_channels: int + spin_width: int @property def state_width(self) -> int: @@ -210,16 +213,7 @@ def spin_slice(self) -> slice: is empty for a spin-free profile. """ stop = self.output_width - 2 - self.degree_channels[0] - vector_width = 1 + 2 * self.spin_channels - width = ( - 0 - if not self.has_spin - else vector_width * (vector_width + 1) // 2 - + 2 - + 2 * self.degree_channels[2] - + 2 * self.spin_channels - ) - return slice(stop - width, stop) + return slice(stop - self.spin_width, stop) def descriptor_profile( @@ -295,6 +289,7 @@ def descriptor_profile( gram_base=gram_base, bispectrum_base=bispectrum_base, spin_channels=spin_channels, + spin_width=spin_dim, ) @@ -1120,7 +1115,7 @@ def _cpu_descriptor( node_count = atype.shape[0] channels = type_embedding.shape[1] type_count = type_embedding.shape[0] - has_spin = spin is not None and spin.numel() != 0 + has_spin = spin is not None and spin.ndim == 2 profile = descriptor_profile(int(channels), int(lmax), has_spin) radial_modes = 0 if pair_mixing.numel() == 0 else int(pair_mixing.shape[2]) @@ -1517,7 +1512,7 @@ def _cpu_forward(*args: Any) -> tuple[torch.Tensor, torch.Tensor]: profile = descriptor_profile( int(args[9].shape[1]), int(args[20]), - args[16].numel() != 0, + args[16].ndim == 2, ) state = torch.zeros( descriptor.shape[0], @@ -1580,7 +1575,7 @@ def _forward_fake( degree_floor, ) profile = descriptor_profile( - int(type_embedding.shape[1]), int(lmax), spin.numel() != 0 + int(type_embedding.shape[1]), int(lmax), spin.ndim == 2 ) descriptor = torch.empty( atype.shape[0], @@ -1605,7 +1600,7 @@ def _backward_fake( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: del descriptor_gradient, state spin = args[15] - has_spin = spin.numel() != 0 + has_spin = spin.ndim == 2 # Each absent output is allocated separately: the schema declares three # unannotated results, so two of them may not share storage. return ( @@ -1636,7 +1631,7 @@ def _cpu_backward( """ del state spin = args[15] - has_spin = spin.numel() != 0 + has_spin = spin.ndim == 2 value = edge_vec.detach().clone().requires_grad_(True) moment = spin.detach().clone().requires_grad_(has_spin) with torch.enable_grad(): @@ -1756,7 +1751,7 @@ def compressed_operator_arguments( Radial table, ordered caches, readout projections, coupling tables, output calibration, and the native spin block. """ - empty = descriptor.compress_spin_type[:0] + empty = descriptor.compress_spin_type.new_empty(0) return ( descriptor.compress_data, descriptor.compress_pair_film, diff --git a/deepmd/kernels/cuda/edge_force_virial.py b/deepmd/kernels/cuda/edge_force_virial.py index eb254bb605..9e8e2ecd93 100644 --- a/deepmd/kernels/cuda/edge_force_virial.py +++ b/deepmd/kernels/cuda/edge_force_virial.py @@ -111,6 +111,13 @@ def frame_scalar_sum( return torch.ops.deepmd.frame_scalar_sum(node_scalar, n_node_per_frame) +def _has_spin_cotangent( + edge_spin_gradient: torch.Tensor, +) -> bool: + """Identify a spin cotangent by its tensor rank.""" + return edge_spin_gradient.ndim == 2 + + def _fake( g_e: torch.Tensor, edge_vec: torch.Tensor, @@ -130,7 +137,9 @@ def _fake( g_e.new_empty(node_capacity, 3), g_e.new_empty(node_capacity if want_atom_virial else 0, 3, 3), g_e.new_empty(n_frame, 3, 3), - g_e.new_empty(node_capacity if edge_spin_gradient.numel() else 0, 3), + g_e.new_empty(node_capacity, 3) + if _has_spin_cotangent(edge_spin_gradient) + else g_e.new_empty(0), ) @@ -151,7 +160,9 @@ def _canonical_fake( g_e.new_empty(node_capacity, 3), g_e.new_empty(node_capacity if want_atom_virial else 0, 3, 3), g_e.new_empty(n_frame, 3, 3), - g_e.new_empty(node_capacity if edge_spin_gradient.numel() else 0, 3), + g_e.new_empty(node_capacity, 3) + if _has_spin_cotangent(edge_spin_gradient) + else g_e.new_empty(0), ) @@ -181,7 +192,7 @@ def _cpu( ) if not want_atom_virial: atom_virial = atom_virial.new_zeros(0, 3, 3) - if edge_spin_gradient.numel(): + if _has_spin_cotangent(edge_spin_gradient): # A masked edge carries no force and no moment, so the two reductions # must agree on which edges exist. contribution = edge_spin_gradient @@ -191,7 +202,7 @@ def _cpu( node_capacity, 3, dtype=g_e.dtype, device=g_e.device ).index_add_(0, edge_index[0], contribution) else: - magnetic_force = g_e.new_zeros(0, 3) + magnetic_force = g_e.new_empty(0) return force, atom_virial, virial, magnetic_force @@ -324,8 +335,8 @@ def edge_force_virial( n_node_per_frame : torch.Tensor Per-frame node counts with shape (nf,), int64. edge_spin_gradient : torch.Tensor - Per-edge magnetic cotangent with shape (E, 3), or an empty tensor when - the model carries no magnetic degree of freedom. + Per-edge magnetic cotangent with shape (E, 3), or a rank-one empty + sentinel when the model carries no magnetic degree of freedom. node_capacity : int Padded node-axis size ``N`` (may be a ``SymInt`` under tracing). want_atom_virial : bool @@ -341,8 +352,8 @@ def edge_force_virial( virial : torch.Tensor Per-frame virial with shape (nf, 3, 3). magnetic_force : torch.Tensor - Per-source total of the magnetic cotangent with shape (N, 3), or an - empty (0, 3) tensor when no spin cotangent was supplied. + Per-source total of the magnetic cotangent with shape (N, 3), or a + rank-one empty sentinel when no spin cotangent was supplied. """ ensure_registered() return torch.ops.deepmd.edge_force_virial( @@ -387,7 +398,8 @@ def canonical_edge_force_virial( n_node_per_frame Per-frame node counts with shape ``(nf,)``. edge_spin_gradient - Per-edge magnetic cotangent with shape ``(E, 3)``, or empty. + Per-edge magnetic cotangent with shape ``(E, 3)``, or a rank-one + empty sentinel when the model carries no magnetic degree of freedom. node_capacity Flat node count ``N``. want_atom_virial @@ -395,10 +407,10 @@ def canonical_edge_force_virial( Returns ------- - tuple[torch.Tensor, torch.Tensor, torch.Tensor] + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] Force, optional atom virial, frame virial, and the per-source total of - the magnetic cotangent, the last empty when no spin cotangent was - supplied. + the magnetic cotangent, the last a rank-one empty sentinel when no + spin cotangent was supplied. """ ensure_registered() return torch.ops.deepmd.canonical_edge_force_virial( diff --git a/deepmd/kernels/cuda/graph_fitting.py b/deepmd/kernels/cuda/graph_fitting.py index 2243dba2e8..c22ea4ad06 100644 --- a/deepmd/kernels/cuda/graph_fitting.py +++ b/deepmd/kernels/cuda/graph_fitting.py @@ -79,9 +79,15 @@ def node_tile() -> int: def op_available() -> bool: - """Whether the C++ ``deepmd::graph_fitting`` op is loaded.""" - op = getattr(torch.ops.deepmd, "graph_fitting", None) - return isinstance(op, torch._ops.OpOverloadPacket) + """Whether every C++ fitting operator used by this module is loaded.""" + operators = ( + getattr(torch.ops.deepmd, "graph_fitting", None), + getattr(torch.ops.deepmd, "graph_fitting_backward", None), + getattr(torch.ops.deepmd, "graph_fitting_energy_gradient", None), + ) + return all( + isinstance(operator, torch._ops.OpOverloadPacket) for operator in operators + ) def fitting_eligible(fit: Any) -> bool: diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 7e10e99466..2c05bae7af 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -31,17 +31,20 @@ ) -def custom_huber_loss( - predictions: torch.Tensor, targets: torch.Tensor, delta: float = 1.0 -) -> torch.Tensor: - error = targets - predictions - abs_error = torch.abs(error) - quadratic_loss = 0.5 * torch.pow(error, 2) +def _huber_from_residual(residual: torch.Tensor, delta: float = 1.0) -> torch.Tensor: + abs_error = torch.abs(residual) + quadratic_loss = 0.5 * torch.pow(residual, 2) linear_loss = delta * (abs_error - 0.5 * delta) loss = torch.where(abs_error <= delta, quadratic_loss, linear_loss) return torch.mean(loss) +def custom_huber_loss( + predictions: torch.Tensor, targets: torch.Tensor, delta: float = 1.0 +) -> torch.Tensor: + return _huber_from_residual(targets - predictions, delta) + + class EnergyStdLoss(TaskLoss): def __init__( self, @@ -426,11 +429,8 @@ def forward( ) # [nf, nloc, 3] huber_ncomp = 3 else: - diff_3 = (force_label - force_pred).reshape( - _nf, _nloc, 3 - ) norm_2d = torch.linalg.vector_norm( - diff_3.reshape(-1, 3), ord=2, dim=1 + diff_f_3d.reshape(-1, 3), ord=2, dim=1 ).reshape(_nf, _nloc) abs_n = norm_2d quad_n = 0.5 * torch.square(norm_2d) @@ -452,21 +452,19 @@ def forward( ) else: if not self.f_use_norm: - l_huber_loss = custom_huber_loss( - force_pred.reshape(-1), - force_label.reshape(-1), + l_huber_loss = _huber_from_residual( + diff_f, delta=self._huber_delta_force, ) else: force_diff_norm = torch.linalg.vector_norm( - (force_label - force_pred).reshape(-1, 3), + diff_f.reshape(-1, 3), ord=2, dim=1, keepdim=True, ) - l_huber_loss = custom_huber_loss( + l_huber_loss = _huber_from_residual( force_diff_norm, - torch.zeros_like(force_diff_norm), delta=self._huber_delta_force, ) loss += pref_f * l_huber_loss @@ -486,9 +484,8 @@ def forward( torch.abs(diff_f_3d), maskf, 3 ) else: - diff_3 = (force_label - force_pred).reshape(_nf, _nloc, 3) norm_2d = torch.linalg.vector_norm( - diff_3.reshape(-1, 3), ord=2, dim=1 + diff_f_3d.reshape(-1, 3), ord=2, dim=1 ).reshape(_nf, _nloc) # One L2 norm per atom, hence one label per atom. l1_f_masked = masked_atom_mean( @@ -500,14 +497,10 @@ def forward( loss += (pref_f * l1_f_masked).to(GLOBAL_PT_FLOAT_PRECISION) else: if not self.f_use_norm: - l1_force_loss = F.l1_loss( - force_label.reshape(-1), - force_pred.reshape(-1), - reduction="mean", - ) + l1_force_loss = torch.mean(torch.abs(diff_f)) else: l1_force_loss = torch.linalg.vector_norm( - (force_label - force_pred).reshape(-1, 3), + diff_f.reshape(-1, 3), ord=2, dim=1, keepdim=True, diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index cb1172c369..c3f3632413 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -329,7 +329,7 @@ def _without_magnetic_force( tuple of torch.Tensor The same outputs followed by an empty magnetic force. """ - return (*output, output[2].new_empty((0, 3))) + return (*output, output[2].new_empty(0)) @BaseDescriptor.register("se_atten") @@ -1097,10 +1097,11 @@ def fused_energy_force_graph( Collapses this descriptor, the energy fitting and the analytic force / virial assembly into one value-returning CUDA operator (no autograd - tape). Returns ``(energy, atom_energy, force, virial, atom_virial)``, or - ``None`` when the descriptor or fitting is not fused-eligible or the - operator library is unavailable -- the caller then uses the autograd - lower. The geo-compressed descriptor dispatches to its tabulated operator + tape). Returns ``(energy, atom_energy, force, virial, atom_virial, + force_mag)``, with a rank-one empty magnetic-force sentinel, or ``None`` + when the descriptor or fitting is not fused-eligible or the operator + library is unavailable -- the caller then uses the autograd lower. The + geo-compressed descriptor dispatches to its tabulated operator (:func:`~deepmd.kernels.cuda.dpa1.graph_compress.dpa1_graph_compress_energy_force`); the embedding-MLP descriptor to :func:`~deepmd.kernels.cuda.dpa1.graph_energy_force.dpa1_graph_energy_force`. @@ -1120,11 +1121,15 @@ def fused_energy_force_graph( Combined fitting and atomic-model bias with shape (ntypes,). do_atomic_virial : bool Whether to also assemble the per-atom virial. + spin : torch.Tensor or None + Optional per-node magnetic moments. DPA1 does not consume them; + supplying one makes the caller select a spin-capable fallback. Returns ------- tuple[torch.Tensor, ...] or None - ``(energy, atom_energy, force, virial, atom_virial)``, or ``None``. + ``(energy, atom_energy, force, virial, atom_virial, force_mag)``, + or ``None``. """ from deepmd.kernels.cuda.graph_fitting import ( fitting_eligible, diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 7aaa9cb8dd..791ce3da42 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -139,7 +139,7 @@ def edge_energy_deriv( source_order, source_row_ptr, n_node, - edge_vec.new_zeros(0, 3), + edge_vec.new_empty(0), n_cap, do_atomic_virial, ) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index e78be39250..28853ecd11 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -245,7 +245,14 @@ def forward_lower_canonical_graph( "virial": virial, "mask": output_mask.to(torch.int32), } - if force_mag is not None and force_mag.numel() != 0: + if spin is not None: + if force_mag is None or force_mag.ndim != 2: + raise RuntimeError( + "canonical native-spin inference did not return a " + "per-node magnetic force" + ) + result["force_mag"] = force_mag + elif force_mag is not None and force_mag.ndim == 2: result["force_mag"] = force_mag if do_atomic_virial: result["atom_virial"] = atom_virial diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 563bd4f1a4..bbd3262342 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -133,7 +133,7 @@ def _fused_energy_force_graph( var + "_derv_c_redu": virial.reshape(nf, 1, 9), "mask": output_mask.to(torch.int32), } - if force_mag.numel() != 0: + if force_mag.ndim == 2: ret[var + "_derv_r_mag"] = force_mag.reshape(n, 1, 3) elif spin is not None: # A moment was supplied but this descriptor emits no magnetic force, diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 3c896fd0d1..7aa95b3c9b 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1404,7 +1404,8 @@ def deserialize_to_file( raise ValueError( "native-spin models implement only the NeighborGraph and compact " f"canonical lowers (got lower_kind={lower_kind!r}); use " - "lower_kind='graph' with a .pt2 output." + "lower_kind='graph', or lower_kind='dpa4c_canonical' for an " + "eligible compressed DPA4C model, with a .pt2 output." ) # A graph lower deploys the fused inference pipeline. The trace runs at # DP_CUDA_INFER >= 2 so the analytic backward and CSR scatter remain custom @@ -2276,6 +2277,18 @@ def _deserialize_to_file_pt2( ) metadata["output_keys"] = output_keys + # A charge-state update swaps constants in one compiled artifact. A model + # that also needs the with-comm lower would own two independent constant + # buffers, so reject that unsupported combination before compiling either + # artifact rather than packaging two states that can silently diverge. + charge_state_descriptor = _charge_state_descriptor(data, metadata) + has_comm_artifact = bool(metadata.get("has_comm_artifact")) + if charge_state_descriptor is not None and has_comm_artifact: + raise ValueError( + "a charge-state fold cannot be packaged with a with-comm lower " + "because each compiled lower owns independent constants" + ) + # On CUDA, aggressive kernel fusion (default realize_opcount_threshold=30) # causes NaN in the backward pass (force/virial) of attention-based # descriptors (DPA1, DPA2). Setting threshold=0 prevents fusion and @@ -2331,7 +2344,6 @@ def _deserialize_to_file_pt2( # Charge-state fold. Present only for a compressed charge-conditioned # descriptor, whose frozen tables the deployment rebuilds once when the # state becomes known. - charge_state_descriptor = _charge_state_descriptor(data, metadata) charge_state_bytes: bytes | None = None if charge_state_descriptor is not None: metadata["charge_state_constants"] = _match_charge_state_constants( @@ -2345,7 +2357,6 @@ def _deserialize_to_file_pt2( # passing extends across rank boundaries. The flag was computed # from the model in ``_collect_metadata`` and is already in # ``metadata`` here. - has_comm_artifact = bool(metadata.get("has_comm_artifact")) with_comm_bytes: bytes | None = None with_comm_output_keys: list[str] | None = None if has_comm_artifact: diff --git a/deepmd/utils/eval_metrics.py b/deepmd/utils/eval_metrics.py index 622f9fad23..2b87014bd5 100644 --- a/deepmd/utils/eval_metrics.py +++ b/deepmd/utils/eval_metrics.py @@ -353,7 +353,7 @@ def compute_full_validation_spin_metrics( The energy term reuses per-atom energy errors. Forces are split into a real-atom term over all atoms and a magnetic term over the magnetic atoms selected by ``mask_mag``. A periodic system additionally reports stress, - the virial divided by the cell volume. + the negated virial divided by the cell volume. Parameters ---------- @@ -495,7 +495,7 @@ class FullValidationMetricProfile: needs_spin=False, log_header_note=( "# E uses per-atom energy, F uses component-wise force errors, " - "and S uses stress, the virial divided by the cell volume.\n" + "and S uses stress, the negated virial divided by the cell volume.\n" ), compute_system_metrics=compute_full_validation_energy_metrics, ) @@ -552,7 +552,7 @@ class FullValidationMetricProfile: log_header_note=( "# E uses per-atom energy, FR uses component-wise real-atom force " "errors, FM uses magnetic-atom force errors, and S uses stress, the " - "virial divided by the cell volume.\n" + "negated virial divided by the cell volume.\n" ), compute_system_metrics=compute_full_validation_spin_metrics, ) diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 78f8c6f1e0..d4cef081b6 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -442,6 +442,8 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, * @param[in] nloc Number of owned local nodes. * @param[in] nall_nodes Total local-plus-halo node count. * @param[in] edge_storage Number of edge storage slots. + * @note API version 29 used signed int64 source and source-order arrays; + * API version 30 uses the uint32 arrays declared here. * @since API version 29 */ extern void DP_DeepPotComputeCanonicalGraphGPU( @@ -1275,7 +1277,7 @@ extern void DP_DeepSpinComputeNListf3(DP_DeepSpin* dp, * @param[in] nloc Number of owned local nodes. * @param[in] nall_nodes Total local-plus-halo node count. * @param[in] edge_storage Number of edge storage slots. - * @since API version 29 + * @since API version 30 */ extern void DP_DeepSpinComputeCanonicalGraphGPU( DP_DeepSpin* dp, @@ -1297,14 +1299,14 @@ extern void DP_DeepSpinComputeCanonicalGraphGPU( /** * @brief Query whether the compact canonical graph ABI is active for a DP spin * model. - * @since API version 29 + * @since API version 30 */ extern bool DP_DeepSpinUsesCanonicalGraphInference(DP_DeepSpin* dp); /** * @brief Query whether a DP spin model is served under the native spin scheme * rather than the virtual-atom scheme. - * @since API version 29 + * @since API version 30 */ extern bool DP_DeepSpinUsesNativeSpinScheme(DP_DeepSpin* dp); diff --git a/source/api_c/include/c_api_internal.h b/source/api_c/include/c_api_internal.h index 63de2ebd35..e871d62b75 100644 --- a/source/api_c/include/c_api_internal.h +++ b/source/api_c/include/c_api_internal.h @@ -1,5 +1,6 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include +#include #include "DataModifier.h" #include "DeepBaseModel.h" @@ -8,6 +9,16 @@ #include "DeepTensor.h" #include "neighbor_list.h" +namespace deepmd { +namespace c_api_internal { + +std::vector copy_charge_spin(const double* charge_spin, + int numb_chg_spin, + int dim_chg_spin); + +} // namespace c_api_internal +} // namespace deepmd + // catch deepmd::deepmd_exception and store it in dp->exception // return nothing #define DP_REQUIRES_OK(dp, xx) \ diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index cee85e258b..c35a4a1763 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -771,6 +771,11 @@ bool validate_model_devi_nframes(DP_DeepBaseModelDevi* dp, const int nframes) { return false; } +} // namespace + +namespace deepmd { +namespace c_api_internal { + /** * @brief Copy a charge/spin condition supplied by a C caller. * @@ -788,10 +793,18 @@ std::vector copy_charge_spin(const double* charge_spin, "the charge/spin condition carries " + std::to_string(numb_chg_spin) + " values but the model expects " + std::to_string(dim_chg_spin)); } + if (numb_chg_spin == 0) { + return {}; + } + if (charge_spin == nullptr) { + throw deepmd::deepmd_exception( + "the charge/spin condition pointer is null for a non-empty input"); + } return std::vector(charge_spin, charge_spin + numb_chg_spin); } -} // namespace +} // namespace c_api_internal +} // namespace deepmd template void DP_DeepPotModelDeviCompute_variant( @@ -2780,8 +2793,9 @@ int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp) { return dp->dp.dim_chg_spin(); } void DP_DeepPotSetChargeSpin(DP_DeepPot* dp, const double* charge_spin, const int numb_chg_spin) { - DP_REQUIRES_OK(dp, dp->dp.set_charge_spin(copy_charge_spin( - charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); + DP_REQUIRES_OK( + dp, dp->dp.set_charge_spin(deepmd::c_api_internal::copy_charge_spin( + charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); } int DP_DeepSpinGetDimChgSpin(DP_DeepSpin* dp) { return dp->dp.dim_chg_spin(); } @@ -2789,8 +2803,9 @@ int DP_DeepSpinGetDimChgSpin(DP_DeepSpin* dp) { return dp->dp.dim_chg_spin(); } void DP_DeepSpinSetChargeSpin(DP_DeepSpin* dp, const double* charge_spin, const int numb_chg_spin) { - DP_REQUIRES_OK(dp, dp->dp.set_charge_spin(copy_charge_spin( - charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); + DP_REQUIRES_OK( + dp, dp->dp.set_charge_spin(deepmd::c_api_internal::copy_charge_spin( + charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); } bool DP_DeepPotIsAParamNAll(DP_DeepPot* dp) { @@ -2836,8 +2851,9 @@ int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp) { void DP_DeepPotModelDeviSetChargeSpin(DP_DeepPotModelDevi* dp, const double* charge_spin, const int numb_chg_spin) { - DP_REQUIRES_OK(dp, dp->dp.set_charge_spin(copy_charge_spin( - charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); + DP_REQUIRES_OK( + dp, dp->dp.set_charge_spin(deepmd::c_api_internal::copy_charge_spin( + charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); } bool DP_DeepPotModelDeviIsAParamNAll(DP_DeepPotModelDevi* dp) { @@ -2922,8 +2938,9 @@ int DP_DeepSpinModelDeviGetDimChgSpin(DP_DeepSpinModelDevi* dp) { void DP_DeepSpinModelDeviSetChargeSpin(DP_DeepSpinModelDevi* dp, const double* charge_spin, const int numb_chg_spin) { - DP_REQUIRES_OK(dp, dp->dp.set_charge_spin(copy_charge_spin( - charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); + DP_REQUIRES_OK( + dp, dp->dp.set_charge_spin(deepmd::c_api_internal::copy_charge_spin( + charge_spin, numb_chg_spin, dp->dp.dim_chg_spin()))); } bool DP_DeepSpinModelDeviIsAParamNAll(DP_DeepSpinModelDevi* dp) { diff --git a/source/api_c/tests/test_deepmd_exception.cc b/source/api_c/tests/test_deepmd_exception.cc index 2c63868d5e..86abc12743 100644 --- a/source/api_c/tests/test_deepmd_exception.cc +++ b/source/api_c/tests/test_deepmd_exception.cc @@ -254,3 +254,9 @@ TEST(TestChargeSpinValidation, rejects_invalid_size) { EXPECT_THROW(deepmd::hpp::validate_charge_spin(charge_spin, 2, 2, tiled), deepmd::hpp::deepmd_exception); } + +TEST(TestChargeSpinValidation, c_boundary_rejects_null_nonempty_input) { + EXPECT_TRUE(deepmd::c_api_internal::copy_charge_spin(nullptr, 0, 0).empty()); + EXPECT_THROW(deepmd::c_api_internal::copy_charge_spin(nullptr, 2, 2), + deepmd::deepmd_exception); +} diff --git a/source/api_cc/include/NativeSpinPTExpt.h b/source/api_cc/include/NativeSpinPTExpt.h index 4830bcd67f..51b0c2dd69 100644 --- a/source/api_cc/include/NativeSpinPTExpt.h +++ b/source/api_cc/include/NativeSpinPTExpt.h @@ -73,7 +73,7 @@ struct CanonicalGraphTensorPack; class NativeSpinPTExpt : public DeepSpinBackend { public: NativeSpinPTExpt(); - virtual ~NativeSpinPTExpt(); + ~NativeSpinPTExpt() override; NativeSpinPTExpt(const std::string& model, const int& gpu_rank = 0, const std::string& file_content = ""); @@ -85,25 +85,25 @@ class NativeSpinPTExpt : public DeepSpinBackend { **/ void init(const std::string& model, const int& gpu_rank = 0, - const std::string& file_content = ""); + const std::string& file_content = "") override; - double cutoff() const { + double cutoff() const override { assert(inited); return rcut; }; - int numb_types() const { + int numb_types() const override { assert(inited); return ntypes; }; - int numb_types_spin() const { + int numb_types_spin() const override { assert(inited); return ntypes_spin; }; - int dim_fparam() const { + int dim_fparam() const override { assert(inited); return dfparam; }; - int dim_aparam() const { + int dim_aparam() const override { assert(inited); return daparam; }; @@ -138,9 +138,9 @@ class NativeSpinPTExpt : public DeepSpinBackend { * @param[in] charge_spin The condition, of length ``dim_chg_spin()``. **/ void set_charge_spin(const std::vector& charge_spin) override; - void get_type_map(std::string& type_map); - bool is_aparam_nall() const { return false; }; - bool has_default_fparam() const { + void get_type_map(std::string& type_map) override; + bool is_aparam_nall() const override { return false; }; + bool has_default_fparam() const override { assert(inited); return has_default_fparam_; }; @@ -161,7 +161,7 @@ class NativeSpinPTExpt : public DeepSpinBackend { const std::vector& box, const std::vector& fparam, const std::vector& aparam, - const bool atomic); + const bool atomic) override; void computew(std::vector& ener, std::vector& force, std::vector& force_mag, @@ -174,7 +174,7 @@ class NativeSpinPTExpt : public DeepSpinBackend { const std::vector& box, const std::vector& fparam, const std::vector& aparam, - const bool atomic); + const bool atomic) override; void computew(std::vector& ener, std::vector& force, std::vector& force_mag, @@ -190,7 +190,7 @@ class NativeSpinPTExpt : public DeepSpinBackend { const int& ago, const std::vector& fparam, const std::vector& aparam, - const bool atomic); + const bool atomic) override; void computew(std::vector& ener, std::vector& force, std::vector& force_mag, @@ -206,7 +206,7 @@ class NativeSpinPTExpt : public DeepSpinBackend { const int& ago, const std::vector& fparam, const std::vector& aparam, - const bool atomic); + const bool atomic) override; // Charge/spin-aware overloads. This backend serves the condition in force // rather than marshalling one per call, so a condition named here is diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 7331eab8b2..a9894e2f04 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -3,6 +3,7 @@ #ifdef BUILD_PYTORCH #include +#include #include #include @@ -417,6 +418,23 @@ struct CanonicalGraphTensorPack { torch::Tensor source_order; }; +/** + * @brief Return the scalar type of compact canonical graph indices. + * + * Unsigned 32-bit tensors entered the public C++ API in PyTorch 2.3. Older + * CPU-only libtorch releases can still build DeePMD-kit, but cannot execute + * the GPU-only compact canonical graph path. + */ +inline at::ScalarType canonicalGraphIndexType() { +#if TORCH_VERSION_MAJOR > 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 3) + return torch::kUInt32; +#else + throw deepmd_exception( + "compact canonical graph inference requires PyTorch 2.3 or later"); +#endif +} + inline CanonicalGraphTensorPack compactCanonicalGraph( const GraphTensorPack& graph) { const std::int64_t edge_count = @@ -428,23 +446,24 @@ inline CanonicalGraphTensorPack compactCanonicalGraph( throw deepmd_exception( "compact canonical graph exceeds the uint32 edge-index range"); } + const auto index_type = canonicalGraphIndexType(); auto source = torch::zeros({storage_count}, - graph.edge_index.options().dtype(torch::kUInt32)); + graph.edge_index.options().dtype(index_type)); auto edge_vec = torch::zeros({storage_count, 3}, graph.edge_vec.options().dtype(torch::kFloat32)); auto source_order = torch::arange(storage_count, graph.edge_index.options().dtype(torch::kInt64)) - .to(torch::kUInt32); + .to(index_type); if (edge_count > 0) { source.slice(0, 0, edge_count) .copy_(graph.edge_index.select(0, 0) .slice(0, 0, edge_count) - .to(torch::kUInt32)); + .to(index_type)); edge_vec.slice(0, 0, edge_count) .copy_(graph.edge_vec.slice(0, 0, edge_count).to(torch::kFloat32)); source_order.slice(0, 0, edge_count) - .copy_(graph.source_order.slice(0, 0, edge_count).to(torch::kUInt32)); + .copy_(graph.source_order.slice(0, 0, edge_count).to(index_type)); } return {graph.atype, graph.n_node, diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index b1792066c8..a2d3eeddcd 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -427,19 +427,17 @@ void DeepPotPTExpt::set_charge_spin(const std::vector& charge_spin) { " values but the model expects " + std::to_string(settable_chgspin)); } - // Route one: the condition of every later forward pass that is not given - // one explicitly. This is the whole mechanism for an uncompressed model, - // which reads the condition as an ordinary input. - default_chg_spin_ = charge_spin; - // Route two: a compressed descriptor has folded the condition into frozen - // tables that the lower holds as constants, so serving another condition - // means rebuilding those tables and writing them over the constants. + // All allocation completes before the compiled constants change. The final + // vector swap is non-throwing, so the visible default follows the installed + // constants without opening a second failure point. + std::vector next_charge_spin(charge_spin); if (charge_state_fold_) { charge_state_fold_->apply(charge_spin, gpu_enabled ? torch::Device(torch::kCUDA, gpu_id) : torch::Device(torch::kCPU), *loader); } + default_chg_spin_.swap(next_charge_spin); } std::vector DeepPotPTExpt::run_model( @@ -822,7 +820,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // non-message-passing, or nghost == 0: the regular path is always safe. // message-passing, multi-rank: requires the with-comm artifact. // message-passing, single-rank: requires the atom map (mapping tensor). - if (has_message_passing_ && nghost > 0) { + if (!lower_input_is_canonical_ && has_message_passing_ && nghost > 0) { if (multi_rank && !has_comm_artifact_) { throw deepmd::deepmd_exception( "Multi-rank LAMMPS .pt2 inference requires the model to be " @@ -2562,8 +2560,9 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( torch::TensorOptions().dtype(torch::kFloat32).device(device); const auto opt_i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); - const auto opt_u32 = - torch::TensorOptions().dtype(torch::kUInt32).device(device); + const auto opt_u32 = torch::TensorOptions() + .dtype(deepmd::canonicalGraphIndexType()) + .device(device); auto atype = torch::from_blob(const_cast(d_atype), {nall_nodes}, opt_i64); auto source = torch::from_blob(const_cast(d_source), diff --git a/source/api_cc/src/DeepSpinPTExpt.cc b/source/api_cc/src/DeepSpinPTExpt.cc index fbbf46c689..dcf2075b19 100644 --- a/source/api_cc/src/DeepSpinPTExpt.cc +++ b/source/api_cc/src/DeepSpinPTExpt.cc @@ -760,13 +760,13 @@ void DeepSpinPTExpt::compute(ENERGYVTYPE& ener, } } - // Folding ghost neighbours onto their local owners reads an owner for every - // extended atom out of the mapping, which the preceding matrix guarantees - // only for a message-passing model. Without a mapping the owner table falls - // back to the identity, whose ghost entries are not local atoms; - // establishing the precondition here reports it once, ahead of any tensor, - // instead of leaving it to the per-edge lookup inside the topology build. - if (!use_with_comm && nghost_real > 0 && !atom_map_present) { + // Edge and graph lowers fold ghost neighbours onto local owners before the + // model runs. Dense lowers consume the extended neighbor list directly and + // retain their established identity-mapping fallback. + const bool folds_ghosts_to_local = + lower_input_is_edge_ || lower_input_is_graph_; + if (folds_ghosts_to_local && !use_with_comm && nghost_real > 0 && + !atom_map_present) { throw deepmd::deepmd_exception( "This .pt2 lower folds ghost neighbours onto their local owners, " "which needs an owner for each of the " + diff --git a/source/api_cc/src/NativeSpinPTExpt.cc b/source/api_cc/src/NativeSpinPTExpt.cc index bfa46bc71e..adf308052e 100644 --- a/source/api_cc/src/NativeSpinPTExpt.cc +++ b/source/api_cc/src/NativeSpinPTExpt.cc @@ -63,6 +63,22 @@ void reject_unsupported_parametric_inputs(const std::vector& fparam, } } +/** + * @brief Validate one Cartesian vector per atom before indexing caller data. + */ +template +void validate_cartesian_input(const std::vector& values, + const std::size_t atom_count, + const char* name) { + const std::size_t expected = atom_count * 3; + if (values.size() != expected) { + throw deepmd::deepmd_exception( + std::string(name) + " holds " + std::to_string(values.size()) + + " values but " + std::to_string(atom_count) + " atoms require " + + std::to_string(expected)); + } +} + /** * @brief Build the frame-parameter input of the conditional graph tail. * @@ -275,11 +291,13 @@ void NativeSpinPTExpt::init(const std::string& model, // The compact lower is traced with the nine graph and moment inputs // alone, so a model declaring any conditioning width has no slot to // receive it. - if (dfparam > 0 || daparam > 0) { + if (dfparam > 0 || daparam > 0 || dchgspin > 0) { throw deepmd::deepmd_exception( - "the compact canonical native-spin ABI has no fparam / aparam slot, " + "the compact canonical native-spin ABI has no fparam, aparam, or " + "charge/spin slot, " "but this model declares dim_fparam=" + std::to_string(dfparam) + ", dim_aparam=" + std::to_string(daparam) + + ", dim_chg_spin=" + std::to_string(dchgspin) + "; freeze it with the graph lower instead."); } } @@ -405,19 +423,17 @@ void NativeSpinPTExpt::set_charge_spin(const std::vector& charge_spin) { " values but the model expects " + std::to_string(settable_chgspin)); } - // Route one: the condition of every later forward pass that is not given - // one explicitly. This is the whole mechanism for an uncompressed model, - // which reads the condition as an ordinary input. - default_chg_spin_ = charge_spin; - // Route two: a compressed descriptor has folded the condition into frozen - // tables that the lower holds as constants, so serving another condition - // means rebuilding those tables and writing them over the constants. + // All allocation completes before the compiled constants change. The final + // vector swap is non-throwing, so the visible default follows the installed + // constants without opening a second failure point. + std::vector next_charge_spin(charge_spin); if (charge_state_fold_) { charge_state_fold_->apply(charge_spin, gpu_enabled ? torch::Device(torch::kCUDA, gpu_id) : torch::Device(torch::kCPU), *loader); } + default_chg_spin_.swap(next_charge_spin); } void NativeSpinPTExpt::get_type_map(std::string& type_map_str) { @@ -546,6 +562,8 @@ void NativeSpinPTExpt::compute(ENERGYVTYPE& ener, std::is_same::value ? torch::kFloat32 : torch::kFloat64; const int nall = static_cast(atype.size()); const int nframes = 1; + validate_cartesian_input(coord, atype.size(), "coord"); + validate_cartesian_input(spin, atype.size(), "spin"); // Drop the atoms whose LAMMPS type maps to NULL: the model never sees them, // and select_map scatters the results back onto the full atom list. @@ -564,7 +582,7 @@ void NativeSpinPTExpt::compute(ENERGYVTYPE& ener, // communication; a single rank folds ghosts onto their local owners, which // needs the LAMMPS atom map to resolve an owner. const bool multi_rank = (lmp_list.nprocs > 1); - if (!multi_rank && nghost > 0 && lmp_list.mapping == nullptr) { + if (!multi_rank && nghost_real > 0 && lmp_list.mapping == nullptr) { throw deepmd::deepmd_exception( "single-rank inference folds ghost neighbours onto their local owners " "through the LAMMPS atom map; add 'atom_modify map yes' to the input, " @@ -794,6 +812,12 @@ void NativeSpinPTExpt::compute(ENERGYVTYPE& ener, std::is_same::value ? torch::kFloat32 : torch::kFloat64; const int nloc = static_cast(atype.size()); const int nframes = 1; + if (atype.empty()) { + throw deepmd::deepmd_exception( + "standalone native-spin inference requires at least one atom"); + } + validate_cartesian_input(coord, atype.size(), "coord"); + validate_cartesian_input(spin, atype.size(), "spin"); // === Step 1. Supply a box when the caller has none === // An isolated cluster is embedded in an orthorhombic cell wide enough that @@ -1169,8 +1193,9 @@ void NativeSpinPTExpt::compute_canonical_graph_gpu( torch::TensorOptions().dtype(torch::kFloat64).device(device); const auto opt_i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); - const auto opt_u32 = - torch::TensorOptions().dtype(torch::kUInt32).device(device); + const auto opt_u32 = torch::TensorOptions() + .dtype(deepmd::canonicalGraphIndexType()) + .device(device); CanonicalGraphTensorPack graph; graph.atype = torch::from_blob(const_cast(d_atype), {nall_nodes}, opt_i64); diff --git a/source/api_cc/src/commonPTExpt.h b/source/api_cc/src/commonPTExpt.h index 5ab6a4f42b..bb67ffd0f5 100644 --- a/source/api_cc/src/commonPTExpt.h +++ b/source/api_cc/src/commonPTExpt.h @@ -838,16 +838,22 @@ class ChargeStateFold { " tables but the archive names " + std::to_string(constants_.size()) + " constants; it cannot serve a runtime charge state"); } - std::unordered_map update; + // The inactive buffer is a complete model image. A partial inactive update + // copies tensor constants and buffers but omits ordinary parameters, so it + // cannot be swapped into service safely. + auto* runner = target.get_runner(); + auto constants = runner->extract_constants_map(/*use_inactive=*/false); for (size_t ii = 0; ii < tables.size(); ++ii) { // An unnamed output belongs to a mechanism this model has disabled and // has no constant to reach. if (!constants_[ii].empty()) { - update[constants_[ii]] = tables[ii]; + constants[constants_[ii]] = tables[ii]; } } - target.update_constant_buffer(update, /*use_inactive=*/false, - /*validate_full_updates=*/false); + target.load_constants(constants, /*use_inactive=*/true, + /*check_full_update=*/true); + runner->swap_constant_buffer(); + runner->free_inactive_constant_buffer(); } private: diff --git a/source/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc index f99ef18673..267edc63dd 100644 --- a/source/api_cc/tests/test_neighbor_list_data.cc +++ b/source/api_cc/tests/test_neighbor_list_data.cc @@ -140,6 +140,10 @@ TEST(TestNeighborListData, RoundTripWithEmptyRows) { #ifdef BUILD_PYTORCH TEST(TestNeighborListData, CompactCanonicalGraphDropsMaskedGuards) { +#if TORCH_VERSION_MAJOR < 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR < 3) + GTEST_SKIP() << "uint32 tensors require PyTorch 2.3 or later"; +#endif GraphTensorPack graph; graph.atype = torch::tensor({0}, torch::kInt64); graph.n_node = torch::tensor({1}, torch::kInt64); @@ -154,12 +158,13 @@ TEST(TestNeighborListData, CompactCanonicalGraphDropsMaskedGuards) { graph.source_row_ptr = torch::tensor({0, 1}, torch::kInt64); const auto compact = compactCanonicalGraph(graph); - EXPECT_EQ(compact.source.scalar_type(), torch::kUInt32); + EXPECT_EQ(compact.source.scalar_type(), deepmd::canonicalGraphIndexType()); EXPECT_EQ(compact.source.numel(), 2); EXPECT_EQ(compact.edge_vec.scalar_type(), torch::kFloat32); EXPECT_EQ(compact.edge_vec.size(0), 2); - EXPECT_TRUE(torch::equal(compact.source_order, - torch::tensor({0, 1}, torch::kUInt32))); + EXPECT_TRUE( + torch::equal(compact.source_order, + torch::tensor({0, 1}, deepmd::canonicalGraphIndexType()))); EXPECT_EQ(compact.destination_row_ptr.select(0, 1).item(), 1); EXPECT_EQ(compact.source_row_ptr.select(0, 1).item(), 1); } diff --git a/source/lmp/compact_canonical_graph_kokkos.h b/source/lmp/compact_canonical_graph_kokkos.h index 8584169420..38389d6bd5 100644 --- a/source/lmp/compact_canonical_graph_kokkos.h +++ b/source/lmp/compact_canonical_graph_kokkos.h @@ -211,7 +211,7 @@ class CompactCanonicalGraphKokkos : protected Pointers { Kokkos::deep_copy(d_destination_row_ptr, std::int64_t{0}); Kokkos::deep_copy(d_source_counts, std::uint32_t{0}); if (node_count_int == 0) { - storage_count = min_storage_edges; + storage_count = 0; return; } auto destination_row_ptr = d_destination_row_ptr; diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index 8ca18f8438..fe02503576 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -520,7 +520,8 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { comm_ptr = &comm_list; } - if ((canonical_graph && nnode_m > 0) || nloc_m > 0 || comm_ptr != nullptr) { + if ((canonical_graph && nnode_m > 0) || + (!canonical_graph && (nloc_m > 0 || comm_ptr != nullptr))) { // Fully device-resident inference: raw device pointers in and out. The // edge buffers are produced on the Kokkos stream and consumed by the model // on PyTorch's stream, and the outputs flow back to the Kokkos scatter, so diff --git a/source/lmp/pair_dpa4spin.cpp b/source/lmp/pair_dpa4spin.cpp index 501a49c869..0a8d6e1ba8 100644 --- a/source/lmp/pair_dpa4spin.cpp +++ b/source/lmp/pair_dpa4spin.cpp @@ -312,6 +312,9 @@ void PairDPA4Spin::settings(int narg, char** arg) { ------------------------------------------------------------------------- */ void PairDPA4Spin::coeff(int narg, char** arg) { + if (narg < 2) { + error->all(FLERR, "Incorrect args for pair coefficients"); + } if (!allocated) { allocate(); } @@ -535,13 +538,14 @@ void PairDPA4Spin::compute(int eflag, int vflag) { // list covers, ghosts included; the spin atom style folds the ghost rows // onto their owners in its reverse communication. Dividing the magnetic // force by hbar / |m| turns the energy gradient with respect to the moment - // into the precession force LAMMPS stores, and leaves an atom whose moment - // vanishes with no precession force at all. + // into the precession force LAMMPS stores. Multiplication by |m| preserves + // that relation while making a zero moment well-defined. + const double force_scale = scale[1][1] * force_unit_cvt_factor; + const double magnetic_force_scale = force_scale / kHBar; for (int ii = 0; ii < nall; ++ii) { for (int dd = 0; dd < 3; ++dd) { - f[ii][dd] += scale[1][1] * dforce[3 * ii + dd] * force_unit_cvt_factor; - fm[ii][dd] += scale[1][1] * dforce_mag[3 * ii + dd] / - (kHBar / sp[ii][3]) * force_unit_cvt_factor; + f[ii][dd] += dforce[3 * ii + dd] * force_scale; + fm[ii][dd] += dforce_mag[3 * ii + dd] * sp[ii][3] * magnetic_force_scale; } } diff --git a/source/op/pt/dpa4c_graph_compress.cu b/source/op/pt/dpa4c_graph_compress.cu index e70f96b9ae..45eb90dc69 100644 --- a/source/op/pt/dpa4c_graph_compress.cu +++ b/source/op/pt/dpa4c_graph_compress.cu @@ -168,6 +168,10 @@ Arguments build_arguments(const Payload& payload, payload.edge_index.scalar_type(), "dpa4c_graph_compress: destination_order dtype must match " "edge_index"); + TORCH_CHECK(payload.canonical || + payload.destination_order.numel() == edge_vec.size(0), + "dpa4c_graph_compress: non-canonical input requires one " + "destination_order entry per edge"); TORCH_CHECK(payload.edge_mask.scalar_type() == torch::kBool, "dpa4c_graph_compress: edge_mask must be bool"); TORCH_CHECK( @@ -187,6 +191,8 @@ Arguments build_arguments(const Payload& payload, } TORCH_CHECK(payload.lmax >= 2 && payload.lmax <= 4, "dpa4c_graph_compress: lmax must be 2, 3, or 4"); + TORCH_CHECK(payload.table_max >= payload.rcut, + "dpa4c_graph_compress: table_max must cover rcut"); // The shared mode cache is sized from a compile-time maximum and the split // spline row assumes an even table width, so the rank set is closed. TORCH_CHECK(radial_modes == 0 || radial_modes == 2 || radial_modes == 4 || @@ -232,14 +238,14 @@ Arguments build_arguments(const Payload& payload, // absolute source indices. Whole-system entry points additionally require // the two to describe the same node axis. TORCH_CHECK(payload.atype.size(0) >= payload.node_begin + node_count, - "dpa4c_graph_compress: destination_row_ptr must have N + 1 " - "entries"); + "dpa4c_graph_compress: atype does not cover the destination " + "node window"); // === Native spin === // The three inputs are present together or not at all. ``spin`` spans the // absolute node axis because neighbour lookups address it with source // indices, and ``spin_type`` packs the four per-type scalars a node reads. - const bool has_spin = payload.spin.numel() != 0; + const bool has_spin = payload.spin.dim() == 2; if (has_spin) { for (const torch::Tensor* tensor : {&payload.spin, &payload.spin_pair, &payload.spin_type}) { @@ -262,8 +268,10 @@ Arguments build_arguments(const Payload& payload, "dpa4c_graph_compress: invalid per-type spin table shape"); } else { TORCH_CHECK( - payload.spin_pair.numel() == 0 && payload.spin_type.numel() == 0, - "dpa4c_graph_compress: spin tables require a spin input"); + payload.spin.dim() == 1 && payload.spin.numel() == 0 && + payload.spin_pair.numel() == 0 && payload.spin_type.numel() == 0, + "dpa4c_graph_compress: absent spin must be a rank-one empty tensor, " + "and spin tables require a spin input"); } Arguments arguments; @@ -375,7 +383,7 @@ std::tuple dpa4c_graph_compress( "different node counts"); const int channels = static_cast(type_embedding.size(1)); const Dimensions widths = - profile_dimensions(channels, static_cast(lmax), spin.numel() != 0); + profile_dimensions(channels, static_cast(lmax), spin.dim() == 2); TORCH_CHECK(edge_vec.is_cuda(), "dpa4c_graph_compress: edge_vec must be a CUDA tensor"); const c10::cuda::CUDAGuard device_guard(edge_vec.device()); @@ -449,7 +457,7 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient, rcut, eps, degree_floor}; - const bool has_spin = spin.numel() != 0; + const bool has_spin = spin.dim() == 2; const long node_count = destination_row_ptr.numel() - 1; // The destination row pointer defines the node axis; the type table must // describe exactly that axis, or the two disagree on how many nodes exist. @@ -483,6 +491,11 @@ dpa4c_graph_compress_backward_impl(torch::Tensor descriptor_gradient, return torch::empty({0}, float_options); }; if (node_count == 0) { + if (has_spin) { + return {torch::zeros_like(edge_vec), + torch::empty({node_count, 3}, float_options), + torch::empty(edge_vec.sizes(), float_options)}; + } return {torch::zeros_like(edge_vec), absent(), absent()}; } auto descriptor_gradient_float = @@ -756,7 +769,7 @@ dpa4c_canonical_compress_energy_gradient(torch::Tensor edge_vec, "dpa4c_canonical_compress_energy_gradient: atype and " "destination_row_ptr describe different node counts"); const int channels = static_cast(type_embedding.size(1)); - const bool has_spin = spin.numel() != 0; + const bool has_spin = spin.dim() == 2; const Dimensions widths = profile_dimensions(channels, static_cast(lmax), has_spin); auto f32 = edge_vec.options().dtype(torch::kFloat32); diff --git a/source/op/pt/edge_force_virial.cu b/source/op/pt/edge_force_virial.cu index c512d4a651..844d46eb56 100644 --- a/source/op/pt/edge_force_virial.cu +++ b/source/op/pt/edge_force_virial.cu @@ -28,6 +28,18 @@ constexpr int kThreads = 256; constexpr int kWarpsPerBlock = kThreads / 32; constexpr int kMaximumVirialPartials = 1024; +bool valid_spin_cotangent(const torch::Tensor& edge_gradient, + const torch::Tensor& edge_spin_gradient) { + if (edge_spin_gradient.dim() == 1) { + return edge_spin_gradient.numel() == 0; + } + return edge_spin_gradient.dim() == 2 && edge_spin_gradient.is_cuda() && + edge_spin_gradient.is_contiguous() && + edge_spin_gradient.device() == edge_gradient.device() && + edge_spin_gradient.sizes() == edge_gradient.sizes() && + edge_spin_gradient.scalar_type() == edge_gradient.scalar_type(); +} + #define FORCE_CHECK_LAUNCH(name) \ do { \ const cudaError_t error = cudaGetLastError(); \ @@ -301,6 +313,7 @@ void launch_force_virial(long node_count, const torch::Tensor& source_row_ptr, const torch::Tensor& frame_row_ptr, const torch::Tensor& edge_spin_gradient, + bool has_spin, torch::Tensor& force, torch::Tensor& node_virial, torch::Tensor& magnetic_force, @@ -309,7 +322,6 @@ void launch_force_virial(long node_count, cudaStream_t stream) { const int node_blocks = static_cast((node_count + kWarpsPerBlock - 1) / kWarpsPerBlock); - const bool has_spin = edge_spin_gradient.numel() != 0; auto assemble = [&](auto spin_tag) { edge_force_virial_kernel <<>>( @@ -320,11 +332,9 @@ void launch_force_virial(long node_count, : nullptr, destination_row_ptr.data_ptr(), source_order.data_ptr(), source_row_ptr.data_ptr(), - edge_spin_gradient.numel() ? edge_spin_gradient.data_ptr() - : nullptr, + has_spin ? edge_spin_gradient.data_ptr() : nullptr, force.data_ptr(), node_virial.data_ptr(), - magnetic_force.numel() ? magnetic_force.data_ptr() - : nullptr); + has_spin ? magnetic_force.data_ptr() : nullptr); }; if (has_spin) { assemble(std::true_type{}); @@ -359,8 +369,9 @@ assemble_force_virial(long node_count, ? atom_virial : torch::empty({node_count, 3, 3}, options); auto virial = torch::zeros({frame_count, 3, 3}, options); - auto magnetic_force = - torch::empty({edge_spin_gradient.numel() ? node_count : 0, 3}, options); + const bool has_spin = edge_spin_gradient.dim() == 2; + auto magnetic_force = has_spin ? torch::empty({node_count, 3}, options) + : torch::empty({0}, options); if (node_count == 0 || frame_count == 0) { return {force, atom_virial, virial, magnetic_force}; } @@ -377,20 +388,23 @@ assemble_force_virial(long node_count, launch_force_virial( node_count, frame_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, frame_row_ptr, edge_spin_gradient, force, - node_virial, magnetic_force, virial_partial, virial, stream); + source_row_ptr, frame_row_ptr, edge_spin_gradient, has_spin, + force, node_virial, magnetic_force, virial_partial, virial, + stream); } else if (source_order.scalar_type() == torch::kUInt32) { launch_force_virial( node_count, frame_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, frame_row_ptr, edge_spin_gradient, force, - node_virial, magnetic_force, virial_partial, virial, stream); + source_row_ptr, frame_row_ptr, edge_spin_gradient, has_spin, + force, node_virial, magnetic_force, virial_partial, virial, + stream); } else { launch_force_virial( node_count, frame_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, frame_row_ptr, edge_spin_gradient, force, - node_virial, magnetic_force, virial_partial, virial, stream); + source_row_ptr, frame_row_ptr, edge_spin_gradient, has_spin, + force, node_virial, magnetic_force, virial_partial, virial, + stream); } }); return {force, atom_virial, virial, magnetic_force}; @@ -485,14 +499,10 @@ edge_force_virial(torch::Tensor edge_gradient, destination_order.scalar_type() == source_order.scalar_type(), "edge_force_virial: destination_order and source_order must have the " "same int32, uint32, or int64 dtype"); - TORCH_CHECK( - edge_spin_gradient.numel() == 0 || - (edge_spin_gradient.is_cuda() && edge_spin_gradient.is_contiguous() && - edge_spin_gradient.device() == edge_gradient.device() && - edge_spin_gradient.sizes() == edge_gradient.sizes() && - edge_spin_gradient.scalar_type() == edge_gradient.scalar_type()), - "edge_force_virial: edge_spin_gradient must be empty or match " - "the gradient in device, layout, shape and dtype"); + TORCH_CHECK(valid_spin_cotangent(edge_gradient, edge_spin_gradient), + "edge_force_virial: edge_spin_gradient must match the gradient " + "in device, layout, shape and dtype, or be a rank-one empty " + "sentinel"); return assemble_force_virial(node_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, n_node_per_frame, @@ -536,13 +546,10 @@ canonical_edge_force_virial(torch::Tensor edge_gradient, "canonical_edge_force_virial: row pointers must have N + 1 " "entries"); TORCH_CHECK( - edge_spin_gradient.numel() == 0 || - (edge_spin_gradient.is_cuda() && edge_spin_gradient.is_contiguous() && - edge_spin_gradient.device() == edge_gradient.device() && - edge_spin_gradient.sizes() == edge_gradient.sizes() && - edge_spin_gradient.scalar_type() == edge_gradient.scalar_type()), - "canonical_edge_force_virial: edge_spin_gradient must be empty " - "or match the gradient in device, layout, shape and dtype"); + valid_spin_cotangent(edge_gradient, edge_spin_gradient), + "canonical_edge_force_virial: edge_spin_gradient must match the " + "gradient in device, layout, shape and dtype, or be a rank-one empty " + "sentinel"); auto edge_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); auto destination_order = torch::empty({0}, source_order.options()); diff --git a/source/op/pt/graph_fitting.cu b/source/op/pt/graph_fitting.cu index a9e38a4d18..b4564784db 100644 --- a/source/op/pt/graph_fitting.cu +++ b/source/op/pt/graph_fitting.cu @@ -283,6 +283,39 @@ FittingLayerPlan fitting_layer_plan(const std::vector& ws) { return plan; } +namespace { + +FittingLayerPlan validate_fitting_forward_inputs( + const char* operation, + const torch::Tensor& x, + const torch::Tensor& atype, + const std::vector& ws, + const torch::Tensor& bias_atom_e) { + TORCH_CHECK(x.dim() == 2 && x.is_cuda() && x.is_contiguous() && + x.scalar_type() == torch::kFloat32, + operation, ": x must be contiguous CUDA fp32 with shape (N, D)"); + TORCH_CHECK(atype.dim() == 1 && atype.size(0) == x.size(0) && + atype.is_cuda() && atype.is_contiguous() && + atype.device() == x.device() && + atype.scalar_type() == torch::kInt64, + operation, + ": atype must be contiguous CUDA int64 with shape (N,) on " + "the device of x"); + const FittingLayerPlan plan = fitting_layer_plan(ws); + TORCH_CHECK( + plan.n_layer > 0 && ws[0].dim() == 2 && ws[0].size(0) == x.size(1), + operation, ": the first fitting weight must match the input width"); + TORCH_CHECK(bias_atom_e.dim() == 1 && bias_atom_e.is_cuda() && + bias_atom_e.is_contiguous() && + bias_atom_e.device() == x.device() && + bias_atom_e.scalar_type() == torch::kFloat64, + operation, + ": bias_atom_e must be contiguous CUDA fp64 on the device of x"); + return plan; +} + +} // namespace + // Evaluate the network over one contiguous run of nodes. Every full-width // tensor is indexed from the run's first node, so the same code serves the // whole node axis and a single tile of it. ``saved`` and ``activation`` are @@ -389,10 +422,10 @@ std::tuple graph_fitting( torch::Tensor b_head, torch::Tensor bias_atom_e, int64_t act) { - TORCH_CHECK(x.is_cuda(), "graph_fitting: x must be a CUDA tensor"); + const FittingLayerPlan plan = validate_fitting_forward_inputs( + "graph_fitting", x, atype, ws, bias_atom_e); const c10::cuda::CUDAGuard device_guard(x.device()); const long n_node = x.size(0); - const FittingLayerPlan plan = fitting_layer_plan(ws); auto f32 = x.options().dtype(torch::kFloat32); auto saved = torch::empty({n_node * plan.saved_width()}, f32); auto e = torch::empty({n_node, 1}, x.options().dtype(torch::kFloat64)); @@ -501,16 +534,11 @@ torch::Tensor graph_fitting_energy_gradient(torch::Tensor x, int64_t act, torch::Tensor seed, int64_t tile) { - TORCH_CHECK( - x.is_cuda() && x.is_contiguous() && x.scalar_type() == torch::kFloat32, - "graph_fitting_energy_gradient: x must be contiguous CUDA fp32"); + const FittingLayerPlan plan = validate_fitting_forward_inputs( + "graph_fitting_energy_gradient", x, atype, ws, bias_atom_e); const c10::cuda::CUDAGuard device_guard(x.device()); const long n_node = x.size(0); const long input_width = x.size(1); - const FittingLayerPlan plan = fitting_layer_plan(ws); - TORCH_CHECK(plan.n_layer > 0 && ws[0].size(0) == input_width, - "graph_fitting_energy_gradient: first weight does not match the " - "descriptor width"); auto f32 = x.options().dtype(torch::kFloat32); auto e = torch::empty({n_node, 1}, x.options().dtype(torch::kFloat64)); if (n_node == 0) { @@ -518,9 +546,10 @@ torch::Tensor graph_fitting_energy_gradient(torch::Tensor x, } auto seed_c = seed.contiguous(); TORCH_CHECK( - seed_c.numel() == n_node && seed_c.scalar_type() == torch::kFloat64, - "graph_fitting_energy_gradient: seed must be fp64 with one " - "entry per node"); + seed_c.numel() == n_node && seed_c.scalar_type() == torch::kFloat64 && + seed_c.is_cuda() && seed_c.device() == x.device(), + "graph_fitting_energy_gradient: seed must be CUDA fp64 with one entry " + "per node on the device of x"); const long run = tile > 0 ? std::min(tile, n_node) : n_node; const int slots = plan.n_layer > 1 ? 2 : 1; diff --git a/source/op/pt/graph_ops.h b/source/op/pt/graph_ops.h index 49fbf40fe4..de847ede1d 100644 --- a/source/op/pt/graph_ops.h +++ b/source/op/pt/graph_ops.h @@ -10,6 +10,7 @@ #pragma once +#include #include #include @@ -166,10 +167,11 @@ void fitting_backward_range(cudaStream_t stream, float* d_x); // Scatter dE/d(edge_vec) into per-node force, per-frame virial and (optional) -// per-node virial. A non-empty ``edge_spin_gradient`` adds the per-source total +// per-node virial. A rank-two ``edge_spin_gradient`` adds the per-source total // of the magnetic cotangent, which shares the source grouping the force -// reduction already walks. Returns (force (N, 3), atom_virial (N, 3, 3) or -// empty, virial (nf, 3, 3), magnetic_force (N, 3) or empty). +// reduction already walks; a rank-one empty tensor denotes its absence. +// Returns (force (N, 3), atom_virial (N, 3, 3) or empty, virial (nf, 3, 3), +// magnetic_force (N, 3) or a rank-one empty sentinel). std::tuple edge_force_virial(torch::Tensor g_e, torch::Tensor edge_vec, diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index e6f9baeae7..98160bee83 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -820,7 +820,13 @@ class TestDPModelEnergyLossForceGradAccum: Covers: mse, mae, huber; plus non-mixed no-op. """ - def _make_loss(self, loss_func="mse", use_huber=False, f_use_norm=False): + def _make_loss( + self, + loss_func="mse", + use_huber=False, + f_use_norm=False, + relative_f=None, + ): return EnergyLoss( starter_learning_rate=1.0, start_pref_e=0.0, @@ -836,6 +842,7 @@ def _make_loss(self, loss_func="mse", use_huber=False, f_use_norm=False): loss_func=loss_func, use_huber=use_huber, f_use_norm=f_use_norm, + relative_f=relative_f, ) def _loss_fn(self, loss_obj, model_pred, label, natoms): @@ -936,6 +943,38 @@ def test_f_use_norm_grad_accum(self, loss_func, use_huber): f_B_hat, ) + @pytest.mark.parametrize("masked", [False, True]) + def test_relative_force_norm_uses_normalized_residual(self, masked): + """Vector-norm losses consume the relative-force residual.""" + relative_f = 1.0 + prediction = np.zeros((1, 2, 3), dtype=np.float64) + label_force = np.array( + [[[3.0, 4.0, 0.0], [0.0, 0.0, 2.0]]], + dtype=np.float64, + ) + mask = np.ones((1, 2), dtype=np.float64) if masked else None + model_pred, label = _full_ener_dicts( + 1, + 2, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=mask, + ) + model_pred["force"] = prediction + label["force"] = label_force + label["find_force"] = 1.0 + loss_obj = self._make_loss( + "mae", + f_use_norm=True, + relative_f=relative_f, + ) + + actual = self._loss_fn(loss_obj, model_pred, label, 2) + label_norm = np.linalg.norm(label_force, axis=-1) + residual_norm = label_norm / (label_norm + relative_f) + expected = residual_norm.mean() + assert np.isclose(actual, expected) + def test_no_op_for_non_mixed(self): """All-ones mask gives same force loss as no mask.""" f = _rnd(NP, 3) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 927c9ba0cd..4d0a21ca60 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -794,7 +794,13 @@ class TestPTEnergyLossForceGradAccum: Covers: mse, mae, huber; plus non-mixed no-op. """ - def _make_loss(self, loss_func="mse", use_huber=False, f_use_norm=False): + def _make_loss( + self, + loss_func="mse", + use_huber=False, + f_use_norm=False, + relative_f=None, + ): return EnergyStdLoss( starter_learning_rate=1.0, start_pref_e=0.0, @@ -806,6 +812,7 @@ def _make_loss(self, loss_func="mse", use_huber=False, f_use_norm=False): loss_func=loss_func, use_huber=use_huber, f_use_norm=f_use_norm, + relative_f=relative_f, ) def _run_invariant(self, loss_obj, f_A, f_A_hat, f_B, f_B_hat): @@ -892,6 +899,32 @@ def test_f_use_norm_grad_accum(self, loss_func, use_huber): f_B_hat, ) + @pytest.mark.parametrize("masked", [False, True]) + def test_relative_force_norm_uses_normalized_residual(self, masked): + """Vector-norm losses consume the relative-force residual.""" + relative_f = 1.0 + prediction = torch.zeros(1, 2, 3, dtype=torch.float64, device="cpu") + label_force = torch.tensor( + [[[3.0, 4.0, 0.0], [0.0, 0.0, 2.0]]], + dtype=torch.float64, + device="cpu", + ) + model_pred = {"force": prediction} + if masked: + model_pred["mask"] = torch.ones(1, 2, dtype=torch.float64, device="cpu") + label = {"force": label_force, "find_force": 1.0} + loss_obj = self._make_loss( + "mae", + f_use_norm=True, + relative_f=relative_f, + ) + + actual = _ener_loss_fn(loss_obj, model_pred, label, 2) + label_norm = torch.linalg.vector_norm(label_force, dim=-1) + residual_norm = label_norm / (label_norm + relative_f) + expected = residual_norm.mean() + torch.testing.assert_close(actual, expected) + def test_no_op_for_non_mixed(self): """All-ones mask gives same force loss as no mask.""" f = _t(NP, 3) diff --git a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py index 8adc39a421..e3bdef7ab8 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py @@ -1313,7 +1313,7 @@ def _assert_parity_vs_separate_ops(self, lmax: int) -> None: graph.source_order, graph.source_row_ptr, graph.n_node, - ev.new_zeros(0, 3), + ev.new_empty(0), n, True, ) @@ -1491,7 +1491,7 @@ def test_fused_energy_uses_owned_nodes_only(self) -> None: graph.source_order, graph.source_row_ptr, graph.n_node, - edge_vec.new_zeros(0, 3), + edge_vec.new_empty(0), n_node, True, ) @@ -1637,7 +1637,7 @@ def _assert_parity_vs_separate_ops(self, lmax: int) -> None: graph.source_order, graph.source_row_ptr, graph.n_node, - ev.new_zeros(0, 3), + ev.new_empty(0), n, True, ) @@ -1971,7 +1971,7 @@ def _fused( src_order, src_row_ptr, n_node, - edge_vec.new_zeros(0, 3), + edge_vec.new_empty(0), total, True, ) @@ -2055,7 +2055,7 @@ def test_compact_canonical_parity(self) -> None: source_order, source_row_ptr, n_node, - edge_vec.new_zeros(0, 3), + edge_vec.new_empty(0), total, True, ) @@ -2066,7 +2066,7 @@ def test_compact_canonical_parity(self) -> None: compact.source_row_ptr, compact.source_order, compact.n_node, - edge_vec.new_zeros(0, 3), + edge_vec.new_empty(0), total, True, ) @@ -2124,11 +2124,11 @@ def assemble(device: str) -> torch.Tensor: src_order, src_row_ptr, n_node, - edge_vec.new_zeros(0, 3), + edge_vec.new_empty(0), total, True, ) - self.assertEqual(empty.numel(), 0) + self.assertEqual(tuple(empty.shape), (0,)) with_spin_force = edge_force_virial( g_e, edge_vec, @@ -2145,6 +2145,31 @@ def assemble(device: str) -> torch.Tensor: )[0] torch.testing.assert_close(force, with_spin_force) + # Rank distinguishes an absent spin cotangent from a present one on an + # empty edge axis. The latter still produces one zero row per node. + empty_vec = edge_vec[:0] + empty_rows = torch.zeros( + total + 1, + dtype=torch.int64, + device=edge_vec.device, + ) + empty_magnetic = edge_force_virial( + empty_vec, + empty_vec, + edge_index[:, :0], + mask[:0], + dst_order[:0], + empty_rows, + src_order[:0], + empty_rows, + n_node, + empty_vec, + total, + True, + )[3] + self.assertEqual(tuple(empty_magnetic.shape), (total, 3)) + torch.testing.assert_close(empty_magnetic, torch.zeros_like(empty_magnetic)) + def test_many_small_frames(self) -> None: """Frame reduction is valid beyond the CUDA grid-y limit.""" frame_count = 8192 diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py index c3bb5e0476..83f1d13417 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py @@ -2,6 +2,9 @@ """Numerical contract of the compressed DPA4C CUDA mega kernel.""" import dataclasses +from collections.abc import ( + Sequence, +) import numpy as np import pytest @@ -14,6 +17,7 @@ ) from deepmd.kernels.cuda.dpa4c.graph_compress import ( _cpu_descriptor, + _cpu_forward, _table_lookup, build_compression_artifacts, build_radial_table, @@ -409,7 +413,7 @@ def test_post_compression_statistics_update_snapshot( def _build_charge_descriptor( channels: int = 8, radial_modes: int = 0, - default_chg_spin: list[float] | None = [2.0, 3.0], + default_chg_spin: Sequence[float] | None = (2.0, 3.0), ) -> DescrptDPA4C: """Return a charge-conditioned descriptor with an active condition head. @@ -428,7 +432,9 @@ def _build_charge_descriptor( precision="float32", seed=17, add_chg_spin_ebd=True, - default_chg_spin=default_chg_spin, + default_chg_spin=( + None if default_chg_spin is None else list(default_chg_spin) + ), ) .cuda() .eval() @@ -592,23 +598,24 @@ def test_supported_surface_parity( @_GPU @pytest.mark.parametrize( - ("index", "replacement", "message"), + ("index", "shape", "dtype", "message"), [ - (6, torch.zeros(9, 4, 2, device="cuda"), "PairFiLM"), - (9, torch.zeros(8, 5, 5, device="cuda"), "invalid readout matrix"), - (10, torch.zeros(1, 8, dtype=torch.int32, device="cuda"), "degree triples"), + (6, (9, 4, 2), torch.float32, "PairFiLM"), + (9, (8, 5, 5), torch.float32, "invalid readout matrix"), + (10, (1, 8), torch.int32, "degree triples"), ], ) def test_operator_rejects_inconsistent_artifacts( index: int, - replacement: torch.Tensor, + shape: tuple[int, ...], + dtype: torch.dtype, message: str, ) -> None: """Device-side shape assumptions are enforced at the operator boundary.""" descriptor = _build_descriptor(8) graph, atype = _build_graph(descriptor, canonical=False) arguments = list(_arguments(descriptor, graph, atype)) - arguments[index] = replacement + arguments[index] = torch.zeros(shape, dtype=dtype, device="cuda") with pytest.raises(RuntimeError, match=message): torch.ops.deepmd.dpa4c_graph_compress(graph.edge_vec, *arguments) @@ -1066,7 +1073,7 @@ def test_fused_energy_force_parity( graph.source_order, graph.source_row_ptr, graph.n_node, - edge_vec.new_zeros(0, 3), + edge_vec.new_empty(0), atype.shape[0], True, ) @@ -1211,6 +1218,7 @@ def _build_spin_descriptor( channels: int, lmax: int = 2, radial_modes: int = 0, + device: str = "cuda", ) -> DescrptDPA4C: """Return a spin-conditioned descriptor with a non-unit reference moment. @@ -1232,7 +1240,7 @@ def _build_spin_descriptor( seed=17, use_spin=[True, False], ) - .cuda() + .to(device) .eval() ) descriptor.spin.set_spin_reference(np.array([1.7, 1.0, 1.0])) @@ -1241,6 +1249,73 @@ def _build_spin_descriptor( return descriptor +def _empty_graph(device: str) -> tuple[NeighborGraph, torch.Tensor]: + """Return a native graph whose node and edge axes are both empty.""" + n_node = torch.zeros(1, dtype=torch.int64, device=device) + edge_index = torch.empty(2, 0, dtype=torch.int64, device=device) + edge_vec = torch.empty(0, 3, dtype=torch.float32, device=device) + edge_order = torch.empty(0, dtype=torch.int64, device=device) + row_pointer = torch.zeros(1, dtype=torch.int64, device=device) + return ( + NeighborGraph( + n_node=n_node, + n_local=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=torch.empty(0, dtype=torch.bool, device=device), + destination_order=edge_order, + destination_row_ptr=row_pointer, + source_order=edge_order, + source_row_ptr=row_pointer, + ), + torch.empty(0, dtype=torch.int64, device=device), + ) + + +def test_empty_native_spin_cpu_profile_preserves_spin_contract() -> None: + """The CPU reference keeps the native-spin state width at zero nodes.""" + descriptor = _build_spin_descriptor(8, device="cpu") + graph, atype = _empty_graph("cpu") + arguments = _with_spin( + _arguments(descriptor, graph, atype), + graph.edge_vec, + ) + + output, state = _cpu_forward(graph.edge_vec, *arguments) + + profile = descriptor_profile(8, 2, True) + assert output.shape == (0, profile.output_width) + assert state.shape == (0, profile.state_width) + + +@_GPU +def test_empty_native_spin_cuda_backward_preserves_spin_contract() -> None: + """The CUDA backward distinguishes present spin from its empty axis.""" + descriptor = _build_spin_descriptor(8) + graph, atype = _empty_graph("cuda") + arguments = _with_spin( + _arguments(descriptor, graph, atype), + graph.edge_vec, + ) + output, state = torch.ops.deepmd.dpa4c_graph_compress( + graph.edge_vec, + *arguments, + ) + + edge_gradient, spin_gradient, edge_spin_gradient = ( + torch.ops.deepmd.dpa4c_graph_compress_backward( + torch.empty_like(output), + state, + graph.edge_vec, + *arguments, + ) + ) + + assert edge_gradient.shape == (0, 3) + assert spin_gradient.shape == (0, 3) + assert edge_spin_gradient.shape == (0, 3) + + @_GPU @pytest.mark.parametrize( ("channels", "lmax", "radial_modes"), diff --git a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py index 279b52cd7a..68394c32f1 100644 --- a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py @@ -20,6 +20,11 @@ build_synthetic_graph_inputs, ) +_GPU = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="CUDA is required for compact canonical graph export", +) + def _config() -> dict: return { @@ -379,6 +384,45 @@ def test_compressed_graph_export(monkeypatch: pytest.MonkeyPatch) -> None: assert metadata["graph_edge_dtype"] == "float32" +def test_charge_state_with_comm_is_rejected_before_compilation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Independent compiled lowers cannot share one mutable charge state.""" + from deepmd.pt_expt.utils import ( + serialization, + ) + + descriptor = object() + monkeypatch.setattr( + serialization, + "_trace_and_export", + lambda *args, **kwargs: ( + object(), + {"has_comm_artifact": True}, + {}, + [], + ), + ) + monkeypatch.setattr( + serialization, + "_charge_state_descriptor", + lambda *args, **kwargs: descriptor, + ) + monkeypatch.setattr( + torch._inductor, + "aoti_compile_and_package", + lambda *args, **kwargs: pytest.fail("compilation must not start"), + ) + + with pytest.raises(ValueError, match="independent constants"): + serialization._deserialize_to_file_pt2( + "unused.pt2", + {}, + lower_kind="graph", + ) + + +@_GPU @pytest.mark.parametrize("channels", [8, 64]) def test_compact_canonical_graph_export( monkeypatch: pytest.MonkeyPatch, From 9df1e97cc2869670ff319e297b66d6f001c32c86 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 14 Aug 2026 23:34:07 +0800 Subject: [PATCH 08/10] fix(dpa4c): align export validation with compact lower support --- deepmd/kernels/cuda/dpa4c/canonical.py | 2 +- source/tests/pt_expt/model/test_dpa4_export.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/deepmd/kernels/cuda/dpa4c/canonical.py b/deepmd/kernels/cuda/dpa4c/canonical.py index ce17c55983..ceae538ca7 100644 --- a/deepmd/kernels/cuda/dpa4c/canonical.py +++ b/deepmd/kernels/cuda/dpa4c/canonical.py @@ -52,7 +52,7 @@ def canonical_model_eligible(model: Any) -> bool: def op_available() -> bool: - """Return whether both compact DPA4C descriptor operators are loaded.""" + """Return whether the complete compact DPA4C operator suite is loaded.""" forward = getattr(torch.ops.deepmd, "dpa4c_canonical_compress", None) backward = getattr( torch.ops.deepmd, diff --git a/source/tests/pt_expt/model/test_dpa4_export.py b/source/tests/pt_expt/model/test_dpa4_export.py index aaa6770a26..1f92e23b80 100644 --- a/source/tests/pt_expt/model/test_dpa4_export.py +++ b/source/tests/pt_expt/model/test_dpa4_export.py @@ -479,7 +479,10 @@ def test_native_spin_nlist_deserialize_rejected(tmp_path) -> None: """ model = _build_native_spin_model_cpu() data = {"model": model.serialize()} - with pytest.raises(ValueError, match="only the NeighborGraph lower"): + with pytest.raises( + ValueError, + match="only the NeighborGraph and compact canonical lowers", + ): deserialize_to_file(str(tmp_path / "spin_nlist.pte"), data, lower_kind="nlist") From 3009cd2d8003cfff283f57713c304a3135202f1f Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 17 Aug 2026 00:42:47 +0800 Subject: [PATCH 09/10] fix(dpa4c): address review on metrics, compression and charge states Full validation reports stress alongside the per-atom virial instead of replacing it, so the released `v:mae`/`v:rmse` selectors keep working. Stress becomes a field of the shared energy-type metrics, which repairs the TF2 and JAX validators that project the shared key map and had silently lost the second-rank column. The five parallel tables of a metric profile collapse into one family declaration, and the log reports whichever of stress and per-atom virial the selected metric names. Compression now rejects a non-empty `exclude_types` instead of emitting an artifact that can never reach the fused kernel, and the charge-state domain travels from the model through the archive metadata to the evaluator, so a condition that addresses no table row is refused on both the folded and the input-tensor path. Descriptors that embed the condition continuously declare no ranges and are unaffected. Also repairs a broken `LOG_COLUMN_ORDER` import that left TF2 full validation unimportable, corrects the `@since` of `DP_DeepPotComputeCanonicalGraphGPU`, fixes the LAMMPS spin example workflow, and restructures the DPA4C manual around deployment, including the Kokkos requirement of the device-resident inference path. --- .../dpmodel/atomic_model/base_atomic_model.py | 4 + .../dpmodel/atomic_model/dp_atomic_model.py | 6 + .../atomic_model/linear_atomic_model.py | 14 + deepmd/dpmodel/descriptor/dpa4c.py | 14 + .../dpmodel/descriptor/dpa4c_nn/__init__.py | 2 + .../descriptor/dpa4c_nn/charge_state.py | 15 +- .../descriptor/make_base_descriptor.py | 11 + deepmd/dpmodel/model/base_model.py | 9 + deepmd/dpmodel/model/make_model.py | 4 + deepmd/dpmodel/train/validation.py | 2 +- deepmd/infer/model_test/ener.py | 19 +- deepmd/jax/train/validation.py | 2 +- deepmd/kernels/cuda/dpa4c/graph_compress.py | 10 +- deepmd/pt_expt/descriptor/dpa4c.py | 5 +- deepmd/pt_expt/infer/charge_state.py | 353 +++++++++++++++ deepmd/pt_expt/infer/deep_eval.py | 261 ++--------- deepmd/pt_expt/train/validation.py | 4 +- deepmd/pt_expt/utils/serialization.py | 3 + deepmd/tf2/train/validation.py | 3 +- deepmd/utils/argcheck.py | 7 +- deepmd/utils/eval_metrics.py | 427 ++++++++++-------- doc/model/dpa4c.md | 291 ++++++------ examples/spin/dpa4c/lmp/README.md | 14 +- examples/spin/dpa4c/lmp/in.lammps | 8 +- source/api_c/include/c_api.h | 6 +- .../pt/model/test_dpa4_dpmodel_parity.py | 35 +- source/tests/pt/test_validation.py | 91 ++++ source/tests/pt_expt/descriptor/test_dpa4c.py | 15 + source/tests/pt_expt/infer/test_deep_eval.py | 80 ++++ 29 files changed, 1116 insertions(+), 599 deletions(-) create mode 100644 deepmd/pt_expt/infer/charge_state.py diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 7050a4a62e..9df9cf21ad 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -254,6 +254,10 @@ def get_default_chg_spin(self) -> list[float] | None: """Get the default charge_spin values.""" return None + def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None: + """Get the row range each charge_spin value indexes, or None.""" + return None + def reinit_atom_exclude( self, exclude_types: list[int] = [], diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index dbafd00b0a..d0e8905223 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -154,6 +154,12 @@ def get_default_chg_spin(self) -> list[float] | None: return self.descriptor.get_default_chg_spin() return None + def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None: + """Get the row range each charge_spin value indexes, or None.""" + if self.add_chg_spin_ebd: + return self.descriptor.get_chg_spin_table_ranges() + return None + def uses_graph_lower(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.uses_graph_lower()) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 6d9a93ac65..4bd06c764e 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -740,6 +740,20 @@ def get_default_chg_spin(self) -> "Array | None": lambda m: m.get_default_chg_spin(), )[1] + def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None: + """The shared table row ranges, if the children agree. + + A condition reaches every consuming child, so it must address the + tables of all of them. Children that disagree share no acceptable + state, which the composition reports as an unconstrained domain + rather than silently enforcing one child's tables on the others. + """ + return self._agreed_default( + self._chg_spin_consumers(), + lambda m: m.get_chg_spin_table_ranges() is not None, + lambda m: m.get_chg_spin_table_ranges(), + )[1] + def has_default_fparam(self) -> bool: """Whether every active child shares one default frame parameter.""" return self._agreed_default( diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index 389dbd27f8..cdffbef95b 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -73,6 +73,7 @@ resolve_swiglu_hidden_width, ) from .dpa4c_nn import ( + CHARGE_STATE_TABLE_RANGES, ChargeStateEmbedding, InvariantReadout, OrderedPairFiLM, @@ -2003,6 +2004,19 @@ def get_default_chg_spin(self) -> list[float] | None: """Return the fallback ``[charge, multiplicity]``, if configured.""" return self.default_chg_spin + def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None: + """Return the row range each value of the frame condition indexes. + + The condition is embedded by gathering one row of the charge table and + one of the multiplicity table, so an acceptable state is a pair of + integers inside these half-open ranges. A folded condition indexes the + same tables at rebuild time, so the ranges hold whether or not the + descriptor is compressed. + """ + if self.charge_spin_embedding is None: + return None + return [tuple(rng) for rng in CHARGE_STATE_TABLE_RANGES] + def has_message_passing_across_ranks(self) -> bool: """Return whether intermediate halo communication is required.""" return False diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py index 3d7bc17fd4..a3610297f4 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py @@ -8,6 +8,7 @@ enumerate_degree_triples, ) from .charge_state import ( + CHARGE_STATE_TABLE_RANGES, ChargeStateEmbedding, canonicalize_charge_spin, validate_charge_state, @@ -33,6 +34,7 @@ ) __all__ = [ + "CHARGE_STATE_TABLE_RANGES", "MAX_ANGULAR_DEGREE", "NEIGHBOR_QUADRUPOLE_CHANNELS", "BispectrumLayout", diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py b/deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py index 79dbe52dd1..d0d8a0bf54 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py @@ -85,6 +85,15 @@ #: Half-open range of representable spin multiplicities. MULTIPLICITY_RANGE = (0, MULTIPLICITY_TABLE_ROWS) +#: Name of each value of a charge state, in order, for diagnostics. +CHARGE_STATE_FIELDS = ("charge", "multiplicity") + +#: Half-open row range addressed by each value of a charge state, in order. +#: A condition is a pair of table row indices, so a host-side boundary that +#: knows these ranges can reject an unaddressable state without knowing which +#: descriptor holds the tables. +CHARGE_STATE_TABLE_RANGES = (CHARGE_RANGE, MULTIPLICITY_RANGE) + def validate_charge_state(charge_spin: Any) -> list[float]: """Check that a frame condition addresses a row of each embedding table. @@ -120,11 +129,11 @@ def validate_charge_state(charge_spin: Any) -> list[float]: ) for value, name, (low, high) in zip( values, - ("charge", "multiplicity"), - (CHARGE_RANGE, MULTIPLICITY_RANGE), + CHARGE_STATE_FIELDS, + CHARGE_STATE_TABLE_RANGES, strict=True, ): - if value != int(value): + if not np.isfinite(value) or value != int(value): raise ValueError(f"The {name} must be an integer, got {value}") if not low <= value < high: raise ValueError( diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index f587cc2906..5c2217ab71 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -113,6 +113,17 @@ def get_default_chg_spin(self) -> Any: """Returns the default charge_spin value, or None.""" return None + def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None: + """Returns the row range each charge_spin value indexes, or None. + + A descriptor that embeds the condition by indexing tables reports + one half-open range per value, which makes every acceptable state + an integer tuple inside those ranges. ``None``, the default, means + the condition enters as a continuous quantity and only its width + is constrained. + """ + return None + def get_geo_compress(self) -> bool: """Return whether geometric tabulated compression is active. diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 0ee9763b1f..43405e38dd 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -125,6 +125,15 @@ def get_default_chg_spin(self) -> list | None: """ return None + def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None: + """Return the row range each charge/spin value indexes, or ``None``. + + ``None`` means the condition is a continuous quantity, so only its + width is constrained; a list means every acceptable value is an + integer inside the matching half-open range. + """ + return None + def get_var_name(self) -> str | None: """Return the fitted property's variable name, or ``None`` if this is not a property model. ``is not None`` is the support diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index bee5b79215..9300eafcb2 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -1163,6 +1163,10 @@ def get_default_chg_spin(self) -> list[float] | None: """Get the default charge_spin values.""" return self.atomic_model.get_default_chg_spin() + def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None: + """Get the row range each charge_spin value indexes, or None.""" + return self.atomic_model.get_chg_spin_table_ranges() + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. diff --git a/deepmd/dpmodel/train/validation.py b/deepmd/dpmodel/train/validation.py index d4a36a216a..f652616be2 100644 --- a/deepmd/dpmodel/train/validation.py +++ b/deepmd/dpmodel/train/validation.py @@ -232,7 +232,7 @@ def __init__( restart_training and self.full_val_file.exists() ) self.table_column_specs = [] - for column_name, metric_key in self.profile.column_order: + for column_name, metric_key in self.profile.columns(self.metric_name): _, metric_unit = format_metric_value_for_table(metric_key, 1.0) header_label = f"{column_name}({metric_unit})" self.table_column_specs.append( diff --git a/deepmd/infer/model_test/ener.py b/deepmd/infer/model_test/ener.py index 985fbdcf72..2fd9b84b02 100644 --- a/deepmd/infer/model_test/ener.py +++ b/deepmd/infer/model_test/ener.py @@ -495,7 +495,11 @@ def evaluate_chunk( "find_virial": find_virial, "energy": test_data["energy"], "force": test_data["force"], - **({"virial": test_data["virial"]} if reports_virial else {}), + **( + {"virial": test_data["virial"], "box": box} + if reports_virial + else {} + ), }, natoms=natoms, has_pbc=data.pbc, @@ -514,17 +518,18 @@ def evaluate_chunk( if shared_metrics.virial is None or shared_metrics.virial_per_atom is None: raise RuntimeError("Virial metrics are unavailable for dp test.") # Stress sigma = -virial / volume, in eV/ų (tensile-positive - # convention). + # convention). The errors come from the shared metrics; the + # per-frame tensors below feed the detail file. + errors.update( + shared_metrics.as_weighted_average_errors( + {"stress": ("mae_s", "rmse_s")} + ) + ) volume = np.abs(np.linalg.det(box.reshape([nframes, 3, 3]))).reshape( [nframes, 1] ) prediction_stress = -virial / volume reference_stress = -test_data["virial"] / volume - errors.update( - compute_error_stat( - prediction_stress, reference_stress - ).as_weighted_average_errors("mae_s", "rmse_s") - ) if dp.has_hessian: errors.update( diff --git a/deepmd/jax/train/validation.py b/deepmd/jax/train/validation.py index 9796b19a69..d611a53324 100644 --- a/deepmd/jax/train/validation.py +++ b/deepmd/jax/train/validation.py @@ -73,7 +73,7 @@ def evaluate_all_systems(self) -> dict[str, float]: aggregated = weighted_average([metric for metric in system_metrics if metric]) return { metric_key: float(aggregated[metric_key]) - for _, metric_key in self.profile.column_order + for metric_key, _, _ in self.table_column_specs if metric_key in aggregated } diff --git a/deepmd/kernels/cuda/dpa4c/graph_compress.py b/deepmd/kernels/cuda/dpa4c/graph_compress.py index d1e460d756..ff1a918ed4 100644 --- a/deepmd/kernels/cuda/dpa4c/graph_compress.py +++ b/deepmd/kernels/cuda/dpa4c/graph_compress.py @@ -613,8 +613,16 @@ def build_compression_artifacts( Raises ------ ValueError - If the descriptor is not an fp32 model with a compiled specialization. + If the descriptor excludes type pairs, or is not an fp32 model with a + compiled specialization. """ + if getattr(descriptor, "exclude_types", None): + raise ValueError( + "DPA4C compressed CUDA has no type-exclusion branch, so a " + "descriptor with a non-empty `exclude_types` cannot reach the " + "fused path. Drop the exclusions, or deploy the uncompressed " + "graph archive." + ) sample_parameter = next(descriptor.parameters()) if sample_parameter.dtype != torch.float32: raise ValueError( diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py index c7d76b2d02..b399a17eaa 100644 --- a/deepmd/pt_expt/descriptor/dpa4c.py +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -564,8 +564,9 @@ def enable_compression( Raises ------ ValueError - If compression is already enabled or the descriptor configuration - has no compiled CUDA specialization. + If compression is already enabled, the descriptor excludes type + pairs, or the descriptor configuration has no compiled CUDA + specialization. """ del min_nbor_dist, table_extrapolate, table_stride_2, check_frequency if self.compress: diff --git a/deepmd/pt_expt/infer/charge_state.py b/deepmd/pt_expt/infer/charge_state.py new file mode 100644 index 0000000000..1a3ab60ee3 --- /dev/null +++ b/deepmd/pt_expt/infer/charge_state.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Host-side handling of the frame charge and spin-multiplicity condition. + +A charge-conditioned model reaches its condition by one of two routes. An +uncompressed model reads it as an ordinary forward input, while a compressed +one folds it into frozen tables at export time and serves a different condition +by rebuilding them. Both routes start from the same user request, so this +module owns the step that turns that request into validated charge states, and +the rebuild that the folded route performs. + +What counts as an acceptable value is a property of the loaded model, not of +this layer. A model that embeds the condition by indexing tables declares the +row range of each value, and a condition outside those ranges is rejected here +because neither the gather nor the compiled kernel bounds-checks the index. A +model that embeds the condition continuously declares no ranges, and only the +width of a request is constrained. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import numpy as np +import torch + +if TYPE_CHECKING: + from collections.abc import ( + Sequence, + ) + + from torch._inductor.package import ( + AOTICompiledModel, + ) + +#: Ranges are declared per value, so a request is described by its position. +_VALUE_LABELS = ("first", "second", "third", "fourth") + + +def _check_against_ranges( + states: np.ndarray, table_ranges: Sequence[Sequence[int]] +) -> None: + """Reject any state that addresses no row of the declared tables.""" + if states.shape[1] != len(table_ranges): + raise ValueError( + f"charge_spin states are {states.shape[1]} values wide, but the " + f"model indexes {len(table_ranges)} tables." + ) + # A non-finite value is not a row index either, so it is reported as the + # integrality failure rather than reaching the range comparison, whose + # message would have to render it. + integral = np.isfinite(states) & (states == np.floor(states)) + for column, (low, high) in enumerate(table_ranges): + label = _VALUE_LABELS[column] if column < len(_VALUE_LABELS) else str(column) + values = states[:, column] + offending = values[~integral[:, column]] + if offending.size: + raise ValueError( + f"The {label} charge_spin value indexes a table row and must " + f"be an integer, got {offending[0]}" + ) + outside = values[(values < low) | (values >= high)] + if outside.size: + raise ValueError( + f"The {label} charge_spin value must lie in [{low}, {high}), " + f"got {outside[0]:.0f}" + ) + + +def charge_states( + charge_spin: Any, + width: int, + table_ranges: Sequence[Sequence[int]] | None = None, +) -> np.ndarray: + """Read a requested condition as validated charge states. + + Parameters + ---------- + charge_spin : Any + Requested condition, of any shape holding a whole number of states. + width : int + Number of values in one charge state. + table_ranges : Sequence[Sequence[int]], optional + Half-open row range of each value, as the model declares it. ``None`` + leaves the values unconstrained beyond the width. + + Returns + ------- + np.ndarray + The requested states with shape ``(n, width)``. + + Raises + ------ + ValueError + If the request does not hold at least one whole state, or names a + state that no declared table row answers. + """ + values = np.asarray(charge_spin, dtype=np.float64).reshape(-1) + if values.size == 0 or values.size % width: + raise ValueError( + f"charge_spin carries {values.size} values, which is not a positive " + f"whole number of {width}-wide charge states." + ) + states = values.reshape(-1, width) + if table_ranges is not None: + _check_against_ranges(states, table_ranges) + return states + + +def charge_states_per_frame( + charge_spin: Any, + nframes: int, + dim_chg_spin: int, + table_ranges: Sequence[Sequence[int]] | None = None, +) -> np.ndarray: + """Read a requested condition as one validated state per frame. + + Parameters + ---------- + charge_spin : Any + Requested condition, holding one state for each frame. + nframes : int + Number of frames the forward covers. + dim_chg_spin : int + Number of values in one charge state. + table_ranges : Sequence[Sequence[int]], optional + Half-open row range of each value, as the model declares it. + + Returns + ------- + np.ndarray + The requested states with shape ``(nframes, dim_chg_spin)``. + + Raises + ------ + ValueError + If the request does not hold exactly one valid state per frame. + """ + states = charge_states(charge_spin, dim_chg_spin, table_ranges) + if states.shape[0] != nframes: + raise ValueError( + f"charge_spin must hold one charge state per frame: expected " + f"{nframes} states of width {dim_chg_spin}, got {states.shape[0]}." + ) + return states + + +def single_charge_state( + charge_spin: Any, + width: int, + table_ranges: Sequence[Sequence[int]] | None = None, +) -> tuple[float, ...]: + """Reduce a requested condition to the one state a folded snapshot serves. + + A folded condition lives in tables that are shared by the whole snapshot, + so it is a property of the loaded model rather than of a frame. A request + that names one state per frame is honoured only when every frame names the + same one. + + Parameters + ---------- + charge_spin : Any + Requested condition, of any shape holding a whole number of states. + width : int + Number of values in one charge state. + table_ranges : Sequence[Sequence[int]], optional + Half-open row range of each value, as the model declares it. + + Returns + ------- + tuple[float, ...] + The requested state, with ``width`` values. + + Raises + ------ + ValueError + If the request does not hold at least one valid state, or holds + several states that are not all equal. + """ + states = charge_states(charge_spin, width, table_ranges) + if not bool((states == states[0]).all()): + raise ValueError( + "This model folds one charge state into its frozen tables and " + "therefore serves a single state at a time, but charge_spin names " + f"{states.shape[0]} states that are not all equal." + ) + return tuple(states[0].tolist()) + + +class ChargeStateFold: + """The rebuild of the frozen tables that carry a compressed charge state. + + A compressed charge-conditioned descriptor evaluates its frame condition + once, when the model is frozen, into a handful of tables. Those tables + reach the compiled lower as module constants, so serving a different + condition means rebuilding them and writing them over those constants + rather than re-evaluating the condition on every step. The archive + therefore ships a second compiled artifact that performs the rebuild, + together with the name of the constant each of its outputs replaces. + + Every lower lifts its constants independently, so the names hold only for + the lower they were resolved against at freeze time. Only a compressed + DPA4C descriptor folds a charge state, and that family never carries + message passing across ranks, so an archive with a fold holds exactly one + lower and the question of a second set of names does not arise. + + Parameters + ---------- + model_file : str + Path to the ``.pt2`` archive. + metadata : dict[str, Any] + Parsed archive metadata. + target : AOTICompiledModel + The lower whose constants carry the condition. + + Attributes + ---------- + width : int + Number of values in a charge state this fold accepts. + + Raises + ------ + ValueError + If the archive declares a fold it cannot supply, or names no width + for a charge state. + """ + + def __init__( + self, + model_file: str, + metadata: dict[str, Any], + target: AOTICompiledModel, + ) -> None: + import tempfile + import zipfile + + from torch._inductor import ( + aoti_load_package, + ) + + from deepmd.pt_expt.utils.serialization import ( + PT2_EXTRA_PREFIX, + ) + + self._constants = [str(name) for name in metadata["charge_state_constants"]] + # The lower reads no condition, so what the model accepts is the state + # the snapshot was frozen against, which is also the layout the rebuild + # consumes. + default_chg_spin = metadata.get("default_chg_spin") + if not default_chg_spin: + raise ValueError( + f"'{model_file}' ships a charge-state fold but names no " + "default_chg_spin, so the width of a charge state is unknown." + ) + self.width = len(default_chg_spin) + + entry = PT2_EXTRA_PREFIX + "charge_state.pt2" + with zipfile.ZipFile(model_file, "r") as zf: + if entry not in zf.namelist(): + raise ValueError( + f"Invalid .pt2 file '{model_file}': it declares " + f"charge_state_constants but carries no '{entry}', so it " + "cannot serve a runtime charge state." + ) + archive = zf.read(entry) + # ``aoti_load_package`` reads a path, so the nested archive is extracted + # to a temporary file that this object owns and releases with itself. + self._archive = tempfile.NamedTemporaryFile(suffix=".pt2") + self._archive.write(archive) + self._archive.flush() + self._runner = aoti_load_package(self._archive.name) + self._target = target + self._applied: tuple[float, ...] | None = None + + @classmethod + def load( + cls, + model_file: str, + metadata: dict[str, Any], + target: AOTICompiledModel, + ) -> ChargeStateFold | None: + """Load the rebuild an archive declares, if it declares one. + + The constant-name field is the archive's claim that the rebuild ships + with it, so an archive that declares the names and cannot supply the + rebuild is malformed and fails here rather than degrading silently. + + Parameters + ---------- + model_file : str + Path to the ``.pt2`` archive. + metadata : dict[str, Any] + Parsed archive metadata. + target : AOTICompiledModel + The lower whose constants carry the condition. + + Returns + ------- + ChargeStateFold or None + The fold, or ``None`` when the archive declares none. + """ + if "charge_state_constants" not in metadata: + return None + return cls(model_file, metadata, target) + + def apply(self, charge_spin: tuple[float, ...]) -> None: + """Rebuild the tables for a condition and write them over the constants. + + Rebuilding is skipped when the condition already applies, so a run that + evaluates many frames at one condition pays for it once. + + Applying a condition overwrites loaded module state and is therefore + not safe to interleave with a forward pass. + + Parameters + ---------- + charge_spin : tuple[float, ...] + The condition, with :attr:`width` values. + + Raises + ------ + RuntimeError + If the rebuild returns a different number of tables than the + archive names constants. + """ + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + if charge_spin == self._applied: + return + # The rebuild consumes the condition in the (1, width) float32 layout + # the inference lower would receive. + tables = self._runner( + torch.tensor([charge_spin], dtype=torch.float32, device=DEVICE) + ) + if len(tables) != len(self._constants): + raise RuntimeError( + f"The charge-state fold returned {len(tables)} tables but the " + f"archive names {len(self._constants)} constants; it cannot " + "serve a runtime charge state." + ) + # An unnamed output belongs to a mechanism this model has disabled and + # has no constant to reach. + self._target.load_constants( + {name: table for name, table in zip(self._constants, tables) if name}, + check_full_update=False, + ) + self._applied = charge_spin diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 90dfb17016..337d679ca3 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -58,6 +58,11 @@ from deepmd.pt.utils.auto_batch_size import ( AutoBatchSize, ) +from deepmd.pt_expt.infer.charge_state import ( + ChargeStateFold, + charge_states_per_frame, + single_charge_state, +) from deepmd.pt_expt.utils.edge_schema import ( edge_schema_from_extended, ) @@ -71,9 +76,6 @@ if TYPE_CHECKING: import ase.neighborlist - from torch._inductor.package import ( - AOTICompiledModel, - ) from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, @@ -125,223 +127,6 @@ def _graph_spin_output_key(odef: "OutputVariableDef") -> str | None: return _GRAPH_CATEGORY_TO_KEY.get(odef.category) -def _reshape_charge_spin( - charge_spin: np.ndarray, nframes: int, dim_chg_spin: int -) -> np.ndarray: - charge_spin_arr = np.asarray(charge_spin) - try: - return charge_spin_arr.reshape(nframes, dim_chg_spin) - except ValueError as err: - raise ValueError( - f"charge_spin must be reshape-compatible with ({nframes}, {dim_chg_spin}), " - f"got shape {charge_spin_arr.shape}." - ) from err - - -def _single_charge_state(charge_spin: np.ndarray, width: int) -> tuple[float, ...]: - """Reduce a requested condition to the one state a folded snapshot serves. - - A folded condition lives in tables that are shared by the whole snapshot, - so it is a property of the loaded model rather than of a frame. A request - that names one state per frame is honoured only when every frame names the - same one. - - Parameters - ---------- - charge_spin : np.ndarray - Requested condition, of any shape holding a whole number of states. - width : int - Number of values in one charge state. - - Returns - ------- - tuple[float, ...] - The requested state, with ``width`` values. - - Raises - ------ - ValueError - If the request does not hold at least one whole state, or holds - several states that are not all equal. - """ - values = np.asarray(charge_spin, dtype=np.float64) - if values.size == 0 or values.size % width: - raise ValueError( - f"charge_spin carries {values.size} values, which is not a positive " - f"whole number of {width}-wide charge states." - ) - states = values.reshape(-1, width) - if not bool((states == states[0]).all()): - raise ValueError( - "This model folds one charge state into its frozen tables and " - "therefore serves a single state at a time, but charge_spin names " - f"{states.shape[0]} states that are not all equal." - ) - return tuple(states[0].tolist()) - - -class _ChargeStateFold: - """The rebuild of the frozen tables that carry a compressed charge state. - - A compressed charge-conditioned descriptor evaluates its frame condition - once, when the model is frozen, into a handful of tables. Those tables - reach the compiled lower as module constants, so serving a different - condition means rebuilding them and writing them over those constants - rather than re-evaluating the condition on every step. The archive - therefore ships a second compiled artifact that performs the rebuild, - together with the name of the constant each of its outputs replaces. - - Every lower lifts its constants independently, so the names hold only for - the lower they were resolved against at freeze time. Only a compressed - DPA4C descriptor folds a charge state, and that family never carries - message passing across ranks, so an archive with a fold holds exactly one - lower and the question of a second set of names does not arise. - - Parameters - ---------- - model_file : str - Path to the ``.pt2`` archive. - metadata : dict[str, Any] - Parsed archive metadata. - target : AOTICompiledModel - The lower whose constants carry the condition. - - Attributes - ---------- - width : int - Number of values in a charge state this fold accepts. - - Raises - ------ - ValueError - If the archive declares a fold it cannot supply, or names no width - for a charge state. - """ - - def __init__( - self, - model_file: str, - metadata: dict[str, Any], - target: "AOTICompiledModel", - ) -> None: - import tempfile - import zipfile - - from torch._inductor import ( - aoti_load_package, - ) - - from deepmd.pt_expt.utils.serialization import ( - PT2_EXTRA_PREFIX, - ) - - self._constants = [str(name) for name in metadata["charge_state_constants"]] - # The lower reads no condition, so what the model accepts is the state - # the snapshot was frozen against, which is also the layout the rebuild - # consumes. - default_chg_spin = metadata.get("default_chg_spin") - if not default_chg_spin: - raise ValueError( - f"'{model_file}' ships a charge-state fold but names no " - "default_chg_spin, so the width of a charge state is unknown." - ) - self.width = len(default_chg_spin) - - entry = PT2_EXTRA_PREFIX + "charge_state.pt2" - with zipfile.ZipFile(model_file, "r") as zf: - if entry not in zf.namelist(): - raise ValueError( - f"Invalid .pt2 file '{model_file}': it declares " - f"charge_state_constants but carries no '{entry}', so it " - "cannot serve a runtime charge state." - ) - archive = zf.read(entry) - # ``aoti_load_package`` reads a path, so the nested archive is extracted - # to a temporary file that this object owns and releases with itself. - self._archive = tempfile.NamedTemporaryFile(suffix=".pt2") - self._archive.write(archive) - self._archive.flush() - self._runner = aoti_load_package(self._archive.name) - self._target = target - self._applied: tuple[float, ...] | None = None - - @classmethod - def load( - cls, - model_file: str, - metadata: dict[str, Any], - target: "AOTICompiledModel", - ) -> "_ChargeStateFold | None": - """Load the rebuild an archive declares, if it declares one. - - The constant-name field is the archive's claim that the rebuild ships - with it, so an archive that declares the names and cannot supply the - rebuild is malformed and fails here rather than degrading silently. - - Parameters - ---------- - model_file : str - Path to the ``.pt2`` archive. - metadata : dict[str, Any] - Parsed archive metadata. - target : AOTICompiledModel - The lower whose constants carry the condition. - - Returns - ------- - _ChargeStateFold or None - The fold, or ``None`` when the archive declares none. - """ - if "charge_state_constants" not in metadata: - return None - return cls(model_file, metadata, target) - - def apply(self, charge_spin: tuple[float, ...]) -> None: - """Rebuild the tables for a condition and write them over the constants. - - Rebuilding is skipped when the condition already applies, so a run that - evaluates many frames at one condition pays for it once. - - Applying a condition overwrites loaded module state and is therefore - not safe to interleave with a forward pass. - - Parameters - ---------- - charge_spin : tuple[float, ...] - The condition, with :attr:`width` values. - - Raises - ------ - RuntimeError - If the rebuild returns a different number of tables than the - archive names constants. - """ - from deepmd.pt_expt.utils.env import ( - DEVICE, - ) - - if charge_spin == self._applied: - return - # The rebuild consumes the condition in the (1, width) float32 layout - # the inference lower would receive. - tables = self._runner( - torch.tensor([charge_spin], dtype=torch.float32, device=DEVICE) - ) - if len(tables) != len(self._constants): - raise RuntimeError( - f"The charge-state fold returned {len(tables)} tables but the " - f"archive names {len(self._constants)} constants; it cannot " - "serve a runtime charge state." - ) - # An unnamed output belongs to a mechanism this model has disabled and - # has no constant to reach. - self._target.load_constants( - {name: table for name, table in zip(self._constants, tables) if name}, - check_full_update=False, - ) - self._applied = charge_spin - - def _warn_legacy_edge_vec(metadata: dict) -> None: """Warn once per model load when an edge_vec-schema artifact is opened. @@ -432,7 +217,7 @@ def __init__( self._is_pt2 = model_file.endswith(".pt2") # Only a compressed ``.pt2`` folds its charge state into constants; a # model that reads the condition as an ordinary input needs no rebuild. - self._charge_state_fold: _ChargeStateFold | None = None + self._charge_state_fold: ChargeStateFold | None = None if self._is_pt2: self._load_pt2(model_file) @@ -713,7 +498,7 @@ def _load_pt2(self, model_file: str) -> None: PyTorch 2.11 ``load_pt2`` loader accepts the archive without the "outdated pt2 file" fallback warning. A compressed charge-conditioned archive carries a second compiled artifact beside the inference lower, - which :class:`_ChargeStateFold` loads to serve a runtime condition. + which :class:`ChargeStateFold` loads to serve a runtime condition. """ import zipfile @@ -756,7 +541,7 @@ def _load_pt2(self, model_file: str) -> None: self._pt2_runner = aoti_load_package(model_file) self.exported_module = None - self._charge_state_fold = _ChargeStateFold.load( + self._charge_state_fold = ChargeStateFold.load( model_file, self.metadata, self._pt2_runner ) @@ -920,6 +705,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "has_chg_spin_ebd": model.has_chg_spin_ebd(), "has_default_chg_spin": model.get_default_chg_spin() is not None, "default_chg_spin": model.get_default_chg_spin(), + "chg_spin_table_ranges": model.get_chg_spin_table_ranges(), "is_spin": self._is_spin, "lower_input_kind": "graph" if use_graph_lower else "nlist", } @@ -1149,6 +935,19 @@ def get_dim_chg_spin(self) -> int: return self._dpmodel.get_dim_chg_spin() return int(self.metadata.get("dim_chg_spin", 0)) + def _chg_spin_table_ranges(self) -> list | None: + """Get the row range each charge/spin value indexes, as declared. + + The domain of a condition belongs to the model that consumes it: one + that gathers table rows accepts only integers inside those rows, while + one that embeds the condition continuously accepts any value. A + compiled archive carries the declaration in its metadata, so the same + rule applies with or without a Python model in hand. + """ + if self._dpmodel is not None: + return self._dpmodel.get_chg_spin_table_ranges() + return self.metadata.get("chg_spin_table_ranges") + def _no_runtime_condition_reason(self) -> str: """Explain why the loaded model serves no runtime charge state. @@ -1208,7 +1007,9 @@ def _apply_charge_state(self, charge_spin: np.ndarray | None = None) -> None: return if charge_spin is None: charge_spin = self.metadata["default_chg_spin"] - fold.apply(_single_charge_state(charge_spin, fold.width)) + fold.apply( + single_charge_state(charge_spin, fold.width, self._chg_spin_table_ranges()) + ) def _make_charge_spin_input( self, nframes: int, charge_spin: np.ndarray | None = None @@ -1253,7 +1054,9 @@ def _make_charge_spin_input( return None if charge_spin is not None: return torch.tensor( - _reshape_charge_spin(charge_spin, nframes, dim_chg_spin), + charge_states_per_frame( + charge_spin, nframes, dim_chg_spin, self._chg_spin_table_ranges() + ), dtype=torch.float64, device=DEVICE, ) @@ -1266,7 +1069,13 @@ def _make_charge_spin_input( if hasattr(default_chg_spin, "cpu"): default_chg_spin = default_chg_spin.cpu().numpy() return ( - torch.tensor(default_chg_spin, dtype=torch.float64, device=DEVICE) + torch.tensor( + single_charge_state( + default_chg_spin, dim_chg_spin, self._chg_spin_table_ranges() + ), + dtype=torch.float64, + device=DEVICE, + ) .view(1, dim_chg_spin) .expand(nframes, -1) .contiguous() diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py index 3c13009400..e1a593ccbc 100644 --- a/deepmd/pt_expt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -279,7 +279,7 @@ def __init__( ) self.auto_batch_size = AutoBatchSize(silent=True) self.table_column_specs = [] - for column_name, metric_key in self.profile.column_order: + for column_name, metric_key in self.profile.columns(self.metric_name): _, metric_unit = format_metric_value_for_table( metric_key, 1.0, self.profile ) @@ -431,7 +431,7 @@ def evaluate_all_systems(self) -> dict[str, float]: aggregated = weighted_average([metric for metric in system_metrics if metric]) return { metric_key: float(aggregated[metric_key]) - for _, metric_key in self.profile.column_order + for metric_key, _, _ in self.table_column_specs if metric_key in aggregated } diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 7aa95b3c9b..2269c47b7c 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1116,6 +1116,9 @@ def _collect_metadata( "has_chg_spin_ebd": model.has_chg_spin_ebd(), "has_default_chg_spin": model.get_default_chg_spin() is not None, "default_chg_spin": _metadata_value_to_json(model.get_default_chg_spin()), + # Present only when the condition indexes tables, which is what lets a + # deployment reject an unaddressable state without a Python model. + "chg_spin_table_ranges": model.get_chg_spin_table_ranges(), "fitting_output_defs": fitting_output_defs, # sel_type enables `DeepEval.get_sel_type()` without a dpmodel # round-trip; required for dipole/polar/wfc models in metadata-only diff --git a/deepmd/tf2/train/validation.py b/deepmd/tf2/train/validation.py index 999a12b925..0c97c108a4 100644 --- a/deepmd/tf2/train/validation.py +++ b/deepmd/tf2/train/validation.py @@ -13,7 +13,6 @@ import numpy as np from deepmd.dpmodel.train.validation import ( - LOG_COLUMN_ORDER, FullValidatorBase, ) from deepmd.tf2.common import ( @@ -76,7 +75,7 @@ def evaluate_all_systems(self) -> dict[str, float]: aggregated = weighted_average([metric for metric in system_metrics if metric]) return { metric_key: float(aggregated[metric_key]) - for _, metric_key in LOG_COLUMN_ORDER + for metric_key, _, _ in self.table_column_specs if metric_key in aggregated } diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 84c0726a20..f19240ab45 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -6074,10 +6074,13 @@ def validating_args() -> Argument: "Metric used to determine the best checkpoint during full validation. " "The string is case-insensitive. For energy training the supported " f"values are {energy_metrics}; for spin-energy training they are " - f"{spin_metrics}. `E` and `V` are per-atom metrics, `F` and `FR` use " + f"{spin_metrics}. `E` and `V` are per-atom metrics, `S` is stress, the " + "negated virial divided by the cell volume, `F` and `FR` use " "component-wise force errors, and `FM` uses magnetic-force errors, " "matching `dp test`. The corresponding loss prefactors must not both " - "be 0." + "be 0; `S` and `V` are two presentations of the virial and both " + "require `start_pref_v` and `limit_pref_v`. The validation log reports " + "whichever of `S` and `V` is selected, defaulting to `S`." ) doc_full_val_file = "The file for writing full validation results only. This file is independent from `training.disp_file`." doc_full_val_start = ( diff --git a/deepmd/utils/eval_metrics.py b/deepmd/utils/eval_metrics.py index 2b87014bd5..efcd4de8c4 100644 --- a/deepmd/utils/eval_metrics.py +++ b/deepmd/utils/eval_metrics.py @@ -18,13 +18,18 @@ Callable, ) -# Full validation reports the second-rank response as stress rather than as -# virial, and stress is not a field of :class:`EnergyTypeEvalMetrics`: it needs -# the cell volume. It is therefore contributed by ``_stress_weighted_errors`` -# instead of being projected here. FULL_VALIDATION_WEIGHTED_METRIC_KEYS = { "energy_per_atom": ("mae_e_per_atom", "rmse_e_per_atom"), "force": ("mae_f", "rmse_f"), + "virial_per_atom": ("mae_v_per_atom", "rmse_v_per_atom"), + "stress": ("mae_s", "rmse_s"), +} +# Spin full validation splits the force term into real and magnetic parts, so +# it projects every shared metric except the plain force. +SPIN_FULL_VALIDATION_WEIGHTED_METRIC_KEYS = { + "energy_per_atom": ("mae_e_per_atom", "rmse_e_per_atom"), + "virial_per_atom": ("mae_v_per_atom", "rmse_v_per_atom"), + "stress": ("mae_s", "rmse_s"), } DP_TEST_WEIGHTED_METRIC_KEYS = { "energy": ("mae_e", "rmse_e"), @@ -80,6 +85,7 @@ class EnergyTypeEvalMetrics: force: ErrorStat | None = None virial: ErrorStat | None = None virial_per_atom: ErrorStat | None = None + stress: ErrorStat | None = None def as_weighted_average_errors( self, @@ -146,6 +152,46 @@ def compute_weighted_error_stat( ) +def _compute_stress_error_stat( + virial_prediction: np.ndarray, + virial_reference: np.ndarray, + box: np.ndarray | None, +) -> ErrorStat | None: + """Compute the stress error of one evaluation dataset. + + Stress is the negated virial divided by the cell volume, the + tensile-positive convention ``dp test`` reports. A frame whose cell is + singular carries no stress and is dropped rather than contributing a + divergent entry. + + Parameters + ---------- + virial_prediction : np.ndarray + Predicted virial with shape (nframes, 9), in eV. + virial_reference : np.ndarray + Reference virial with shape (nframes, 9), in eV. + box : np.ndarray | None + Cell vectors with shape (nframes, 9) in Angstrom, or ``None`` when the + caller supplies no cell. + + Returns + ------- + ErrorStat | None + Stress error in eV/Angstrom^3, or ``None`` when no cell is supplied or + no frame has a non-singular cell. + """ + if box is None: + return None + volume = np.abs(np.linalg.det(np.asarray(box).reshape(-1, 3, 3))) + finite = volume > 0.0 + if not np.any(finite): + return None + scale = -1.0 / volume[finite, None] + return compute_error_stat( + virial_prediction[finite] * scale, virial_reference[finite] * scale + ) + + def compute_energy_type_metrics( prediction: dict[str, np.ndarray], test_data: dict[str, np.ndarray], @@ -158,6 +204,7 @@ def compute_energy_type_metrics( force = None virial = None virial_per_atom = None + stress = None if bool(test_data.get("find_energy", 0.0)): energy = compute_error_stat( @@ -177,14 +224,14 @@ def compute_energy_type_metrics( ) if has_pbc and bool(test_data.get("find_virial", 0.0)): - virial = compute_error_stat( - prediction["virial"].reshape(-1, 9), - test_data["virial"].reshape(-1, 9), - ) + virial_prediction = prediction["virial"].reshape(-1, 9) + virial_reference = test_data["virial"].reshape(-1, 9) + virial = compute_error_stat(virial_prediction, virial_reference) virial_per_atom = compute_error_stat( - prediction["virial"].reshape(-1, 9), - test_data["virial"].reshape(-1, 9), - scale=1.0 / natoms, + virial_prediction, virial_reference, scale=1.0 / natoms + ) + stress = _compute_stress_error_stat( + virial_prediction, virial_reference, test_data.get("box") ) return EnergyTypeEvalMetrics( @@ -193,53 +240,10 @@ def compute_energy_type_metrics( force=force, virial=virial, virial_per_atom=virial_per_atom, + stress=stress, ) -def _stress_weighted_errors( - prediction: dict[str, np.ndarray], - test_data: dict[str, np.ndarray], - has_pbc: bool, -) -> dict[str, tuple[float, float]]: - """Return the weighted stress errors of one system. - - Stress is the negated virial divided by the cell volume, the - tensile-positive convention ``dp test`` reports. A frame whose cell is - singular carries no stress and is dropped rather than producing a - divergent entry. - - Parameters - ---------- - prediction : dict[str, np.ndarray] - Model predictions containing ``virial`` with shape ``(nframes, 9)``. - test_data : dict[str, np.ndarray] - Reference labels containing ``virial`` and ``box``, the latter with - shape ``(nframes, 9)`` in Angstrom. - has_pbc : bool - Whether the system is periodic, gating the metric. - - Returns - ------- - dict[str, tuple[float, float]] - Weighted-average-ready stress errors in eV/Angstrom^3, empty when the - system is aperiodic, carries no virial label, or has no frame with a - non-singular cell. - """ - if not (has_pbc and bool(test_data.get("find_virial", 0.0))): - return {} - box = np.asarray(test_data["box"]).reshape(-1, 3, 3) - volume = np.abs(np.linalg.det(box)) - finite = volume > 0.0 - if not np.any(finite): - return {} - scale = -1.0 / volume[finite] - stress = compute_error_stat( - prediction["virial"].reshape(-1, 9)[finite] * scale[:, None], - test_data["virial"].reshape(-1, 9)[finite] * scale[:, None], - ) - return stress.as_weighted_average_errors("mae_s", "rmse_s") - - def compute_spin_force_metrics( force_real_prediction: np.ndarray, force_real_reference: np.ndarray, @@ -329,7 +333,7 @@ def compute_full_validation_energy_metrics( natoms : int The number of atoms per frame, used for per-atom normalization. has_pbc : bool - Whether the system is periodic, gating the virial metrics. + Whether the system is periodic, gating the second-rank metrics. Returns ------- @@ -337,9 +341,7 @@ def compute_full_validation_energy_metrics( Weighted-average-ready ``(value, weight)`` pairs keyed by metric. """ metrics = compute_energy_type_metrics(prediction, test_data, natoms, has_pbc) - errors = metrics.as_weighted_average_errors(FULL_VALIDATION_WEIGHTED_METRIC_KEYS) - errors.update(_stress_weighted_errors(prediction, test_data, has_pbc)) - return errors + return metrics.as_weighted_average_errors(FULL_VALIDATION_WEIGHTED_METRIC_KEYS) def compute_full_validation_spin_metrics( @@ -350,10 +352,9 @@ def compute_full_validation_spin_metrics( ) -> dict[str, tuple[float, float]]: """Compute spin-energy full validation metrics for one system. - The energy term reuses per-atom energy errors. Forces are split into a - real-atom term over all atoms and a magnetic term over the magnetic atoms - selected by ``mask_mag``. A periodic system additionally reports stress, - the negated virial divided by the cell volume. + The energy and second-rank terms come from the shared energy-type metrics. + Forces replace the shared plain-force term with a real-atom term over all + atoms and a magnetic term over the magnetic atoms selected by ``mask_mag``. Parameters ---------- @@ -365,71 +366,75 @@ def compute_full_validation_spin_metrics( natoms : int The number of atoms per frame, used for per-atom normalization. has_pbc : bool - Whether the system is periodic, gating the stress metric. + Whether the system is periodic, gating the second-rank metrics. Returns ------- dict[str, tuple[float, float]] Weighted-average-ready ``(value, weight)`` pairs keyed by metric. """ - errors: dict[str, tuple[float, float]] = {} - if bool(test_data.get("find_energy", 0.0)): - energy_per_atom = compute_error_stat( - prediction["energy"].reshape(-1, 1), - test_data["energy"].reshape(-1, 1), - scale=1.0 / natoms, - ) - errors.update( - energy_per_atom.as_weighted_average_errors( - "mae_e_per_atom", "rmse_e_per_atom" - ) - ) + metrics = compute_energy_type_metrics(prediction, test_data, natoms, has_pbc) + errors = metrics.as_weighted_average_errors( + SPIN_FULL_VALIDATION_WEIGHTED_METRIC_KEYS + ) if bool(test_data.get("find_force", 0.0)): spin_metrics = _spin_force_metrics_from_prediction(prediction, test_data) errors.update( spin_metrics.as_weighted_average_errors(DP_TEST_SPIN_WEIGHTED_METRIC_KEYS) ) - errors.update(_stress_weighted_errors(prediction, test_data, has_pbc)) return errors +@dataclass(frozen=True) +class MetricFamily: + """One quantity a full validation profile reports, as MAE and RMSE. + + Families sharing a loss prefactor pair are alternative presentations of the + same trained quantity, such as the second-rank response reported either as + stress or as per-atom virial. The log table shows exactly one of them. + + Attributes + ---------- + token : str + Family identifier used in ``validation_metric``, such as ``"e"``. + mae_key : str + Internal metric key carrying the mean absolute error. + rmse_key : str + Internal metric key carrying the root mean square error. + unit : tuple[str, float] + Display unit and the factor converting an internal value into it. + prefactors : tuple[str, str] + Loss prefactor keys that must both be active for the family to be + trainable. + """ + + token: str + mae_key: str + rmse_key: str + unit: tuple[str, float] + prefactors: tuple[str, str] + + def metrics(self) -> tuple[tuple[str, str], ...]: + """Return the ``(kind, metric_key)`` pairs this family contributes.""" + return (("mae", self.mae_key), ("rmse", self.rmse_key)) + + @dataclass(frozen=True) class FullValidationMetricProfile: """Metric family definition for one full validation model class. Bundles every aspect that differs between energy-type and spin-energy full validation so the validator stays data-driven instead of branching on the - model class: - - - ``column_order`` defines the ``val.log`` table layout as - ``(header_label, metric_key)`` pairs. - - ``metric_key_map`` maps a normalized ``validation_metric`` token (such as - ``"e:mae"``) to an internal metric key (such as ``"mae_e_per_atom"``). - - ``metric_family_by_key`` maps an internal metric key back to its family, - used for display-unit lookup. - - ``unit_by_family`` maps a family to its ``(display_unit, scale)``. - - ``prefactor_by_metric`` maps a metric token to the loss prefactor keys - that must both be active for the metric to be trainable. - - ``needs_spin`` indicates whether the model consumes a spin input and - emits magnetic forces. - - ``log_header_note`` is the one-line table legend written to ``val.log``. - - ``compute_system_metrics`` turns one system's prediction and reference - into weighted ``(value, weight)`` metric pairs. + model class. Every selectable metric, display unit and loss prefactor pair + derives from ``families``, so each quantity is declared exactly once. Attributes ---------- name : str Profile identifier, either ``"energy"`` or ``"spin"``. - column_order : tuple[tuple[str, str], ...] - Ordered ``(header_label, metric_key)`` pairs for the log table. - metric_key_map : dict[str, str] - Normalized metric token to internal metric key. - metric_family_by_key : dict[str, str] - Internal metric key to family identifier. - unit_by_family : dict[str, tuple[str, float]] - Family identifier to ``(display_unit, scale)``. - prefactor_by_metric : dict[str, tuple[str, str]] - Normalized metric token to ``(start_pref_key, limit_pref_key)``. + families : tuple[MetricFamily, ...] + Reported quantities in table order. Where several families share a + loss prefactor pair, the first one is the default presentation. needs_spin : bool Whether the profile requires spin input and magnetic-force outputs. log_header_note : str @@ -440,11 +445,7 @@ class FullValidationMetricProfile: """ name: str - column_order: tuple[tuple[str, str], ...] - metric_key_map: dict[str, str] - metric_family_by_key: dict[str, str] - unit_by_family: dict[str, tuple[str, float]] - prefactor_by_metric: dict[str, tuple[str, str]] + families: tuple[MetricFamily, ...] needs_spin: bool log_header_note: str compute_system_metrics: Callable[ @@ -452,107 +453,147 @@ class FullValidationMetricProfile: dict[str, tuple[float, float]], ] + @property + def metric_key_map(self) -> dict[str, str]: + """Map a normalized metric token to its internal metric key.""" + return { + f"{family.token}:{kind}": key + for family in self.families + for kind, key in family.metrics() + } + + @property + def metric_family_by_key(self) -> dict[str, str]: + """Map an internal metric key back to its family identifier.""" + return { + key: family.token for family in self.families for _, key in family.metrics() + } + + @property + def unit_by_family(self) -> dict[str, tuple[str, float]]: + """Map a family identifier to its ``(display_unit, scale)``.""" + return {family.token: family.unit for family in self.families} + + @property + def prefactor_by_metric(self) -> dict[str, tuple[str, str]]: + """Map a normalized metric token to its loss prefactor keys.""" + return { + f"{family.token}:{kind}": family.prefactors + for family in self.families + for kind, _ in family.metrics() + } + + def columns(self, metric: str) -> tuple[tuple[str, str], ...]: + """Return the ``val.log`` table layout for a selected metric. + + The table carries one column pair per trained quantity. Where several + families present the same quantity, the selected family wins and the + first declared family is the fallback, so no quantity is reported + twice and the selected metric is always available to the + best-checkpoint selector. + + Parameters + ---------- + metric : str + Normalized ``validation_metric`` token, such as ``"v:rmse"``. + + Returns + ------- + tuple[tuple[str, str], ...] + Ordered ``(header_label, metric_key)`` pairs. + """ + selected = metric.split(":")[0] + shown: dict[tuple[str, str], MetricFamily] = {} + for family in self.families: + if family.prefactors not in shown or family.token == selected: + shown[family.prefactors] = family + return tuple( + (f"{family.token.upper()}_{kind.upper()}", key) + for family in shown.values() + for kind, key in family.metrics() + ) + + +#: Per-atom energy, reported by every profile. +_ENERGY_FAMILY = MetricFamily( + token="e", + mae_key="mae_e_per_atom", + rmse_key="rmse_e_per_atom", + unit=("meV/atom", 1000.0), + prefactors=("start_pref_e", "limit_pref_e"), +) +#: The second-rank response as stress, the default presentation. Declared +#: ahead of the per-atom virial so the table shows stress unless +#: ``validation_metric`` selects the virial instead. +_STRESS_FAMILY = MetricFamily( + token="s", + mae_key="mae_s", + rmse_key="rmse_s", + unit=("meV/ų", 1000.0), + prefactors=("start_pref_v", "limit_pref_v"), +) +#: The second-rank response as virial normalized by the atom count. +_VIRIAL_FAMILY = MetricFamily( + token="v", + mae_key="mae_v_per_atom", + rmse_key="rmse_v_per_atom", + unit=("meV/atom", 1000.0), + prefactors=("start_pref_v", "limit_pref_v"), +) +#: Legend fragment shared by every profile that reports the second-rank term. +_SECOND_RANK_NOTE = ( + "the second-rank column is S, stress as the negated virial divided by the " + "cell volume, or V, virial normalized by natoms, following " + "`validation_metric`.\n" +) ENERGY_FULL_VALIDATION_PROFILE = FullValidationMetricProfile( name="energy", - column_order=( - ("E_MAE", "mae_e_per_atom"), - ("E_RMSE", "rmse_e_per_atom"), - ("F_MAE", "mae_f"), - ("F_RMSE", "rmse_f"), - ("S_MAE", "mae_s"), - ("S_RMSE", "rmse_s"), + families=( + _ENERGY_FAMILY, + MetricFamily( + token="f", + mae_key="mae_f", + rmse_key="rmse_f", + unit=("meV/Å", 1000.0), + prefactors=("start_pref_f", "limit_pref_f"), + ), + _STRESS_FAMILY, + _VIRIAL_FAMILY, ), - metric_key_map={ - "e:mae": "mae_e_per_atom", - "e:rmse": "rmse_e_per_atom", - "f:mae": "mae_f", - "f:rmse": "rmse_f", - "s:mae": "mae_s", - "s:rmse": "rmse_s", - }, - metric_family_by_key={ - "mae_e_per_atom": "e", - "rmse_e_per_atom": "e", - "mae_f": "f", - "rmse_f": "f", - "mae_s": "s", - "rmse_s": "s", - }, - unit_by_family={ - "e": ("meV/atom", 1000.0), - "f": ("meV/Å", 1000.0), - "s": ("meV/ų", 1000.0), - }, - prefactor_by_metric={ - "e:mae": ("start_pref_e", "limit_pref_e"), - "e:rmse": ("start_pref_e", "limit_pref_e"), - "f:mae": ("start_pref_f", "limit_pref_f"), - "f:rmse": ("start_pref_f", "limit_pref_f"), - "s:mae": ("start_pref_v", "limit_pref_v"), - "s:rmse": ("start_pref_v", "limit_pref_v"), - }, needs_spin=False, log_header_note=( - "# E uses per-atom energy, F uses component-wise force errors, " - "and S uses stress, the negated virial divided by the cell volume.\n" + "# E uses per-atom energy, F uses component-wise force errors, and " + + _SECOND_RANK_NOTE ), compute_system_metrics=compute_full_validation_energy_metrics, ) SPIN_FULL_VALIDATION_PROFILE = FullValidationMetricProfile( name="spin", - column_order=( - ("E_MAE", "mae_e_per_atom"), - ("E_RMSE", "rmse_e_per_atom"), - ("FR_MAE", "mae_fr"), - ("FR_RMSE", "rmse_fr"), - ("FM_MAE", "mae_fm"), - ("FM_RMSE", "rmse_fm"), - ("S_MAE", "mae_s"), - ("S_RMSE", "rmse_s"), + families=( + _ENERGY_FAMILY, + MetricFamily( + token="fr", + mae_key="mae_fr", + rmse_key="rmse_fr", + unit=("meV/Å", 1000.0), + prefactors=("start_pref_fr", "limit_pref_fr"), + ), + MetricFamily( + token="fm", + mae_key="mae_fm", + rmse_key="rmse_fm", + unit=("meV/μB", 1000.0), + prefactors=("start_pref_fm", "limit_pref_fm"), + ), + _STRESS_FAMILY, + _VIRIAL_FAMILY, ), - metric_key_map={ - "e:mae": "mae_e_per_atom", - "e:rmse": "rmse_e_per_atom", - "fr:mae": "mae_fr", - "fr:rmse": "rmse_fr", - "fm:mae": "mae_fm", - "fm:rmse": "rmse_fm", - "s:mae": "mae_s", - "s:rmse": "rmse_s", - }, - metric_family_by_key={ - "mae_e_per_atom": "e", - "rmse_e_per_atom": "e", - "mae_fr": "fr", - "rmse_fr": "fr", - "mae_fm": "fm", - "rmse_fm": "fm", - "mae_s": "s", - "rmse_s": "s", - }, - unit_by_family={ - "e": ("meV/atom", 1000.0), - "fr": ("meV/Å", 1000.0), - "fm": ("meV/μB", 1000.0), - "s": ("meV/ų", 1000.0), - }, - prefactor_by_metric={ - "e:mae": ("start_pref_e", "limit_pref_e"), - "e:rmse": ("start_pref_e", "limit_pref_e"), - "fr:mae": ("start_pref_fr", "limit_pref_fr"), - "fr:rmse": ("start_pref_fr", "limit_pref_fr"), - "fm:mae": ("start_pref_fm", "limit_pref_fm"), - "fm:rmse": ("start_pref_fm", "limit_pref_fm"), - "s:mae": ("start_pref_v", "limit_pref_v"), - "s:rmse": ("start_pref_v", "limit_pref_v"), - }, needs_spin=True, log_header_note=( "# E uses per-atom energy, FR uses component-wise real-atom force " - "errors, FM uses magnetic-atom force errors, and S uses stress, the " - "negated virial divided by the cell volume.\n" + "errors, FM uses magnetic-atom force errors, and " + _SECOND_RANK_NOTE ), compute_system_metrics=compute_full_validation_spin_metrics, ) diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md index e93faf03b8..e75f286b54 100644 --- a/doc/model/dpa4c.md +++ b/doc/model/dpa4c.md @@ -8,8 +8,11 @@ family. Where DPA4/SeZM targets the accuracy frontier through equivariant message passing, DPA4C targets the throughput frontier: it reads each local environment once, keeps no message-passing state, and admits a compressed CUDA inference path in which its radial functions are replaced by tabulated splines. -It is intended for large-scale molecular dynamics and as a distillation student -of a DPA4 teacher. + +Choose DPA4C when the run is limited by simulation speed or system size rather +than by the last increment of accuracy: large-scale molecular dynamics, long +trajectories, and distillation from a DPA4 teacher. Choose DPA4 when accuracy +is the binding constraint. DPA4C is selected as a descriptor, `descriptor.type: "dpa4c"`, and pairs with the standard energy fitting network. There is no separate `model.type` scaffold. @@ -25,7 +28,7 @@ dp --pt-expt train input.json copy and adapt. See [training energy models](train-energy.md) for the general workflow shared by all energy models. -## Overview +## How it works DPA4C predicts atomic energies and obtains forces and virials by differentiating the energy, the same conservative formulation used by every @@ -40,18 +43,18 @@ degree {ref}`lmax `, contracts them into rotationally invariant scalars, and passes only those scalars to the fitting network. The neighbor shell is read exactly once: there is no message passing, so an atom's descriptor depends only on the atoms within -{ref}`rcut ` of it. This one-hop -locality is what keeps the per-step cost low and makes the compressed inference -path possible. +{ref}`rcut ` of it. -Two properties follow from the construction and matter in practice. The radial -map is exactly zero at and beyond `rcut`, with continuous derivatives, so the -potential energy surface stays smooth as neighbors cross the cutoff. And every -per-atom quantity is bounded analytically, which is why compression needs -neither an extrapolation region nor overflow checking. +Three consequences shape how the model is used in practice. -If you want the design details, see -[Architecture details](#architecture-details) at the end of this page. +- **One-hop locality** keeps the per-step cost low and removes the cross-rank + halo exchange of intermediate features that a message-passing model needs, so + domain decomposition follows the ordinary pair-style path. +- **Exact smoothness at the cutoff.** The radial map is exactly zero at and + beyond `rcut`, with continuous derivatives, so the potential energy surface + stays smooth as neighbors cross the cutoff. +- **Analytic bounds on every per-atom quantity**, which is why compression needs + neither an extrapolation region nor overflow checking. ## Configuration @@ -93,7 +96,7 @@ DPA4C defaults to `float32` what the compressed CUDA path requires. Double precision is neither necessary nor supported for compressed inference. -### Main options +### Options that matter Every option, with its default and full description, is listed in the {ref}`argument reference `. Four of them @@ -121,11 +124,11 @@ carry the accuracy–cost trade-off: analytic basis that feeds the radial network. > [!IMPORTANT] -> Compressed inference is compiled for -> `radial_modes` in `{0, 2, 4, 8}` only. A model trained with any other value -> trains and runs correctly on the portable path, but `dp --pt-expt compress` -> will reject it. Choose the value with compression in mind if you intend to -> deploy the compressed model. +> The compressed CUDA path is compiled for `channels` in `{8, 16, 32, 64, 128}`, +> `lmax` in `{2, 3, 4}`, and `radial_modes` in `{0, 2, 4, 8}` only. A model +> trained outside those sets trains and evaluates correctly, but +> `dp --pt-expt compress` rejects it. Choose these values with deployment in +> mind. ### Recommended configurations @@ -182,12 +185,14 @@ precision, and the reverse. ## Model compression -Compression replaces the analytic radial functions and their type-pair -modulation with tabulated splines evaluated by fused CUDA kernels. Because the -radial map is analytically bounded and vanishes at `rcut`, the table needs no -extrapolation region and no overflow checking. +Compression is the deployment step. It replaces the analytic radial functions +and their type-pair modulation with tabulated splines evaluated by fused CUDA +kernels, and re-exports the model in the compact canonical graph form that the +fast inference path consumes. Because the radial map is analytically bounded and +vanishes at `rcut`, the table needs no extrapolation region and no overflow +checking. -The workflow is the standard three steps: +Train, freeze, then compress the frozen archive: ```bash dp --pt-expt train input.json @@ -195,25 +200,98 @@ dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph dp --pt-expt compress -i frozen_model.pt2 -o compressed_model.pt2 ``` +The two archives are not interchangeable. `frozen_model.pt2` carries the plain +graph lower and is the uncompressed intermediate; `compressed_model.pt2` carries +the compact canonical graph lower and is what you deploy. Compression selects +that lower on its own, so it takes no lower-kind option of its own. + Only `-s, --step` applies to DPA4C; it sets the uniform spline spacing in Å, and -a smaller value means a finer table and a larger model. -The `--extrapolate`, `--frequency`, and `--training-script` options exist for -descriptors whose tables need a second region, an overflow guard, or a minimum -neighbor distance computed from data; DPA4C needs none of them and ignores them. +a smaller value means a finer table and a larger model. The `--extrapolate`, +`--frequency`, and `--training-script` options exist for descriptors whose +tables need a second region, an overflow guard, or a minimum neighbor distance +computed from data; DPA4C needs none of them and ignores them. Compression requires: - the PyTorch Exportable backend on CUDA; - `precision: "float32"`; -- `channels` in `{8, 16, 32, 64, 128}`, `lmax` in `{2, 3, 4}`, and - `radial_modes` in `{0, 2, 4, 8}`; +- `channels`, `lmax` and `radial_modes` inside the compiled sets listed above; - an empty {ref}`exclude_types `, since - the fused kernel has no type-exclusion branch. A compressed model with - excluded pairs falls back to the portable path. + the fused kernel has no type-exclusion branch. + +`dp --pt-expt compress` reports an explicit error when any of these is not met. +A model that excludes type pairs still trains and runs; deploy it as the +uncompressed graph archive. + +## Running in LAMMPS + +DPA4C uses the PyTorch `.pt2` (AOTInductor) export path and is served by the +`deepmd` pair style: + +```lammps +pair_style deepmd compressed_model.pt2 +pair_coeff * * O H +``` + +### Choosing a pair style + +The compact canonical graph form exists so that the whole step can stay on the +device. Only the Kokkos pair styles use that device-resident entry point; the +host styles run the same archive through a per-step host round trip. Reaching +DPA4C's advertised throughput therefore takes three things together: a +Kokkos-enabled LAMMPS build on the GPU backend, the compressed archive, and +`DP_CUDA_INFER` set at export time as described under +[Inference settings](#inference-settings). + +| Pair style | Build | Accepted archive | Execution | +| ------------- | ------------------------ | ------------------------- | ---------------------------------------------- | +| `deepmd` | any | graph lower or compressed | host round trip each step | +| `deepmd/kk` | Kokkos, GPU backend only | graph lower or compressed | device-resident; compressed uses fused kernels | +| `dpa4spin` | any, `atom_style spin` | graph lower or compressed | host round trip each step | +| `dpa4spin/kk` | Kokkos, GPU backend only | compressed only | device-resident | + +Run under Kokkos with one GPU: + +```bash +lmp -k on g 1 -sf kk -in in.lammps +``` -`dp --pt-expt compress` reports an explicit error when the configuration falls -outside these sets. +### Multiple GPUs + +Because DPA4C performs no message passing, it needs no cross-rank halo exchange +of intermediate features, and MPI domain decomposition follows the ordinary +pair-style path. Launch one MPI rank per GPU and make every target device +visible: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 mpirun -np 4 lmp -in in.lammps +``` + +Use a non-zero neighbor skin, for example `neighbor 2.0 bin`, to keep per-step +GPU memory stable; a zero skin rebuilds the neighbor list every step. + +## Inference settings + +Inference behavior is controlled by environment variables read when the model is +constructed: + +| Environment variable | Default | Effect | +| -------------------- | ------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `DP_CUDA_INFER` | `0` | Fused CUDA kernel level: `0` off, `1` fused descriptor and fitting, `2` additionally fuses force and virial assembly. Levels 1 and 2 are numerically identical. | +| `DP_AMP_INFER` | off | bf16 autocast over the per-edge stage during inference. Independent of the training-time `use_amp`. | +| `DP_TF32_INFER` | `0` | float32 matmul precision: `0` highest, `1` high, `2` medium. | + +A compressed model needs `DP_CUDA_INFER` of at least `1` to reach its fused +path; at `0` it evaluates through the portable path and the compression brings +no speedup. For molecular dynamics sensitive to the smoothness of the potential +energy surface, keep `DP_TF32_INFER=0` and `DP_AMP_INFER=0`. + +> [!IMPORTANT] +> Set these variables **before** running `dp --pt-expt freeze` or +> `dp --pt-expt compress`. The exported `.pt2` is an AOTInductor artifact, so the +> kernel level and precision policy are captured into the graph at export time +> and are **not** re-evaluated when the `.pt2` is later loaded by LAMMPS. ## Native spin @@ -228,7 +306,7 @@ spin gradient of the same energy that yields the conservative force, \mathbf{F}^{m}_i = -\frac{\partial E}{\partial \mathbf{s}_i} . ``` -### Symmetry +### What the descriptor represents The magnetic moment is an axial vector: it is even under spatial inversion and odd under time reversal, whereas a displacement is odd under inversion and even @@ -248,17 +326,18 @@ contracted against one another and against the geometric moments: | Quadrupole | $\sum_j \varphi_c(r_{ij})\,B_2(\hat{\mathbf{s}}_j)$ | Biquadratic exchange, single-ion anisotropy | | Magnitude and magnetic coordination | $\sum_j \varphi_c(r_{ij})\,\lvert\mathbf{s}_j\rvert^2$ and the gated neighbor count | Longitudinal and stoichiometric terms | -The width of the spin block follows the degree-two width of the geometric -descriptor, so it is set by {ref}`channels ` -and has no knob of its own. - Two-body Heisenberg exchange, biquadratic exchange and single-ion anisotropy are represented exactly rather than approximately: each corresponds to a single emitted invariant times a learned radial profile. The Dzyaloshinskii-Moriya interaction is not representable at any order, because the invariant read-out contains no antisymmetric contraction. -### Configuration +The width of the spin block follows the degree-two width of the geometric +descriptor, so it is set by +{ref}`channels ` and has no knob of +its own. + +### Enabling native spin Native spin is requested at the model level, not on the descriptor. The `use_spin` list marks the magnetic types, either as booleans over the type map @@ -301,22 +380,12 @@ unity. A model that declares a magnetic type but receives no moment is rejected rather than evaluated at zero, since the latter is indistinguishable from a broken data pipeline and reports a vanishing magnetic force. -### Running in LAMMPS +### Running a spin model in LAMMPS -A native-spin model is served by the `dpa4spin` pair style, and by -`dpa4spin/kk` under Kokkos. Both require `atom_style spin` and a model frozen -with the compact canonical graph lower, which compression selects on its own -for an eligible DPA4C. Freeze first and compress the frozen artifact, as for -any other DPA4C model: - -```bash -dp --pt-expt freeze -c model.ckpt.pt -o frozen_model -dp --pt-expt compress -i frozen_model.pt2 -o compressed_model.pt2 -``` - -The `lower_input_kind` of `compressed_model.pt2` reads `dpa4c_canonical`, which -is what the pair styles require; `frozen_model.pt2` alone carries the plain -graph lower and is refused. +Freeze and compress exactly as for any other DPA4C model, then select a spin +pair style from the table under +[Choosing a pair style](#choosing-a-pair-style). Both spin styles require +`atom_style spin`: ```lammps atom_style spin @@ -324,73 +393,18 @@ pair_style dpa4spin compressed_model.pt2 pair_coeff * * Ni O ``` -The Kokkos style keeps the graph and the moment in device memory for the whole -step. Ghost moments are supplied by the forward communication that -`atom_style spin` already performs, and the magnetic force is reduced back onto -owning atoms alongside the conservative force, so domain decomposition needs no -additional exchange: - -```bash -lmp -k on g 1 -sf kk -in in.lammps -``` - -A worked example, a rocksalt NiO cell in its type-II antiferromagnetic order, -is provided in `examples/spin/dpa4c/lmp/`. +Ghost moments are supplied by the forward communication that `atom_style spin` +already performs, and the magnetic force is reduced back onto owning atoms +alongside the conservative force, so domain decomposition needs no additional +exchange. `min_style spin` reads the magnetic force from the pair style and relaxes the moment directions. Spin dynamics through `fix nve/spin` requires a LAMMPS build whose fix recognizes this pair style, because the stock fix accumulates the magnetic force only from pair styles matching its own name pattern. -## Export and running in LAMMPS - -DPA4C uses the PyTorch `.pt2` (AOTInductor) export path. Freeze with the graph -lower, which is the form the C++ graph path consumes: - -```bash -dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph -``` - -Use the frozen or compressed `.pt2` with the `deepmd` pair style: - -```lammps -pair_style deepmd compressed_model.pt2 -pair_coeff * * O H -``` - -Because DPA4C performs no message passing, it needs no cross-rank halo exchange -of intermediate features, and MPI domain decomposition follows the ordinary -pair-style path. Launch one MPI rank per GPU and make every target device -visible: - -```bash -CUDA_VISIBLE_DEVICES=0,1,2,3 mpirun -np 4 lmp -in in.lammps -``` - -Use a non-zero neighbor skin, for example `neighbor 2.0 bin`, to keep per-step -GPU memory stable; a zero skin rebuilds the neighbor list every step. - -## Inference settings - -Inference behavior is controlled by environment variables read when the model is -constructed: - -| Environment variable | Default | Effect | -| -------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| `DP_CUDA_INFER` | `0` | Fused CUDA kernel level. `0` disables them. `1` uses the fused descriptor and fitting operators with the force from autograd. `2` additionally collapses descriptor, fitting, and force/virial assembly into one operator, numerically identical to level 1. | -| `DP_AMP_INFER` | off | bf16 autocast over the per-edge stage during inference. Independent of the training-time `use_amp`. | -| `DP_TF32_INFER` | `0` | float32 matmul precision: `0` highest, `1` high, `2` medium. | - -A compressed model requires `DP_CUDA_INFER` of at least `1` to reach its fused -path; at `0` it evaluates through the portable path and the compression brings -no speedup. For molecular dynamics sensitive to the smoothness of the potential -energy surface, keep `DP_TF32_INFER=0` and `DP_AMP_INFER=0`. - -> [!IMPORTANT] -> Set these variables **before** running `dp --pt-expt freeze` or -> `dp --pt-expt compress`. The exported `.pt2` is an AOTInductor artifact, so the -> kernel level and precision policy are captured into the graph at export time -> and are **not** re-evaluated when the `.pt2` is later loaded by LAMMPS. +A worked example, a rocksalt NiO cell in its type-II antiferromagnetic order, +is provided in `examples/spin/dpa4c/lmp/`. ## Data format @@ -400,12 +414,25 @@ DPA4C consumes a mixed-type neighbor list, so it supports both the order consistent across the dataset, the input file, and any downstream `pair_coeff` mapping. -## Architecture details +## Limitations -Optional background on how the descriptor works, linking each part to the -options that control it. Skip it unless you are tuning those options. +- DPA4C is implemented for the PyTorch Exportable backend (`dp --pt-expt`). +- Export uses `.pt2` (AOTInductor); the TorchScript freeze path is not used. +- Model compression requires CUDA, `float32`, and a configuration inside the + compiled sets listed under [Model compression](#model-compression). +- The device-resident inference path requires a Kokkos-enabled LAMMPS build on + the GPU backend. +- The descriptor is one-hop local by construction. Interactions beyond `rcut` + are not represented, and unlike a message-passing model the effective range + cannot be extended by adding layers. +- Native spin requires the `native` scheme; the virtual-atom `deepspin` scheme + is not supported. The Dzyaloshinskii-Moriya interaction is not representable, + as explained under [Native spin](#native-spin). + +## Architecture details -### Edge features +Background on how the descriptor works, linking each part to the options that +control it. Skip it unless you are tuning those options. For every neighbor pair within `rcut`, the interatomic distance is expanded on an analytic radial basis (`basis_type`, with `n_radial` functions) and passed @@ -420,8 +447,6 @@ Each amplitude is multiplied by a smooth cutoff envelope whose value and first derivatives vanish at `rcut`, and by the real spherical harmonics of the neighbor direction up to degree `lmax`. -### Degree-wise moments and invariant read-out - The per-atom state is the sum of these edge contributions, held separately for each angular degree. Degree zero carries `channels` scalar values; higher degrees carry progressively fewer channels, each with `2l + 1` angular @@ -440,25 +465,3 @@ reach the fitting network: Because the contraction is exactly rotationally invariant, the descriptor and hence the energy are invariant under global rotation, and the forces obtained by differentiation are equivariant. - -### Output calibration - -Descriptor statistics are used once, at initialization, to record a fixed -per-coordinate scale that puts the invariant outputs on a comparable footing -before they enter the fitting network. This is an initialization preconditioner, -not a running normalization: no sample-dependent statistic is evaluated during -training or inference, so the model remains a pure function of the atomic -positions and types. - -## Limitations - -- DPA4C is implemented for the PyTorch Exportable backend (`dp --pt-expt`). -- Export uses `.pt2` (AOTInductor); the TorchScript freeze path is not used. -- Model compression requires CUDA, `float32`, and a configuration inside the - compiled sets listed under [Model compression](#model-compression). -- The descriptor is one-hop local by construction. Interactions beyond `rcut` - are not represented, and unlike a message-passing model the effective range - cannot be extended by adding layers. -- Native spin requires the `native` scheme; the virtual-atom `deepspin` scheme - is not supported. The Dzyaloshinskii-Moriya interaction is not representable, - as explained under [Native spin](#native-spin). diff --git a/examples/spin/dpa4c/lmp/README.md b/examples/spin/dpa4c/lmp/README.md index cf7fe1ef64..aa4153f4f6 100644 --- a/examples/spin/dpa4c/lmp/README.md +++ b/examples/spin/dpa4c/lmp/README.md @@ -16,16 +16,20 @@ scheme, see `examples/spin/lmp`; for DPA4 / SeZM native spin, see ## Usage -Train with the configuration in `../input.json`, compress, and freeze. The -pair style requires the compact canonical graph lower, which `--lower-kind auto` selects for a compressed DPA4C; the archive is target-specific and is not -shipped, so freeze locally: +Train with the configuration in `../input.json` and freeze to a `.pt2` archive. +The archive is target-specific and is not shipped, so freeze locally: ```bash dp --pt-expt train ../input.json -dp --pt-expt compress -i model.ckpt.pt -o compressed.pt -dp --pt-expt freeze -c compressed.pt -o frozen_model --lower-kind auto +dp --pt-expt freeze -c model.ckpt.pt -o frozen_model --lower-kind graph ``` +The pair style accepts either a graph-lower archive, as frozen above, or the +compact canonical graph archive that `dp --pt-expt compress` produces, which +runs the fused CUDA kernels. See +[model compression](../../../../doc/model/dpa4c.md#model-compression) for that +workflow, and point `pair_style` at the compressed archive if you follow it. + Run on the host: ```bash diff --git a/examples/spin/dpa4c/lmp/in.lammps b/examples/spin/dpa4c/lmp/in.lammps index a2e39f97c5..050ce53903 100644 --- a/examples/spin/dpa4c/lmp/in.lammps +++ b/examples/spin/dpa4c/lmp/in.lammps @@ -4,12 +4,8 @@ # magnetic force is the negative spin gradient of the energy, so no virtual # atoms are created and the atom count equals the number of physical atoms. # -# The pair style requires a model frozen with the compact canonical graph -# lower, which is what `dp --pt-expt freeze --lower-kind auto` selects for a -# compressed DPA4C: -# -# dp --pt-expt compress -i model.ckpt.pt -o compressed.pt -# dp --pt-expt freeze -c compressed.pt -o frozen_model --lower-kind auto +# The pair style reads a graph-lower archive or the compact canonical graph +# archive produced by compression. See README.md for how to obtain one. # # Substitute `dpa4spin/kk` for the device-resident Kokkos path, which keeps the # graph and the moment in device memory: diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index d4cef081b6..d98d744387 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -442,9 +442,9 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, * @param[in] nloc Number of owned local nodes. * @param[in] nall_nodes Total local-plus-halo node count. * @param[in] edge_storage Number of edge storage slots. - * @note API version 29 used signed int64 source and source-order arrays; - * API version 30 uses the uint32 arrays declared here. - * @since API version 29 + * @note The source and source-order arrays are uint32 as of API version 30; + * API versions 28 and 29 declared them as signed int64. + * @since API version 28 */ extern void DP_DeepPotComputeCanonicalGraphGPU( DP_DeepPot* dp, diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py index dab9fd6833..f35c5d40f5 100644 --- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py +++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py @@ -427,7 +427,8 @@ def test_radial_basis(self, basis_type, exponent) -> None: assert_parity(dp_mod.call(r), pt_mod(to_pt(r))) @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"]) # both bases - def test_radial_basis_roundtrip(self, basis_type) -> None: + @pytest.mark.parametrize("apply_envelope", [True, False]) # both envelope modes + def test_radial_basis_roundtrip(self, basis_type, apply_envelope) -> None: from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( RadialBasis as DPRadialBasis, ) @@ -438,13 +439,45 @@ def test_radial_basis_roundtrip(self, basis_type) -> None: n_radial=12, precision="float64", exponent=7, + apply_envelope=apply_envelope, ) dp_mod2 = DPRadialBasis.deserialize(dp_mod.serialize()) + assert dp_mod2.apply_envelope is apply_envelope r = self._r_grid() np.testing.assert_array_equal( np.asarray(dp_mod.call(r)), np.asarray(dp_mod2.call(r)) ) + @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"]) # both bases + def test_radial_basis_envelope_modes(self, basis_type) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + RadialBasis as DPRadialBasis, + ) + + def make(apply_envelope: bool) -> DPRadialBasis: + return DPRadialBasis( + rcut=self.rcut, + basis_type=basis_type, + n_radial=12, + precision="float64", + exponent=7, + apply_envelope=apply_envelope, + ) + + enveloped, raw = make(True), make(False) + r = self._r_grid() + enveloped_value = np.asarray(enveloped.call(r)) + raw_value = np.asarray(raw.call(r)) + # the two modes differ by exactly one envelope factor + np.testing.assert_array_equal( + enveloped_value, raw_value * np.asarray(enveloped.envelope(r)) + ) + # only the enveloped mode is truncated; consumers of the raw mode owe + # the basis one cutoff factor of their own + outside = r[:, 0] > self.rcut + np.testing.assert_array_equal(enveloped_value[outside], 0.0) + assert np.any(raw_value[outside] != 0.0) + @pytest.mark.parametrize( "mlp_layers", [[16, 32, 24], [16, 24]], diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py index 52b33a9792..c72bc530ed 100644 --- a/source/tests/pt/test_validation.py +++ b/source/tests/pt/test_validation.py @@ -36,13 +36,16 @@ resolve_full_validation_start_step, ) from deepmd.utils.argcheck import ( + is_valid_full_validation_metric, normalize, ) from deepmd.utils.data import ( DataRequirementItem, ) from deepmd.utils.eval_metrics import ( + FULL_VALIDATION_PROFILES, SPIN_FULL_VALIDATION_PROFILE, + compute_full_validation_energy_metrics, compute_full_validation_spin_metrics, ) @@ -769,6 +772,94 @@ def test_spin_profile_splits_real_and_magnetic_forces(self) -> None: self.assertAlmostEqual(metrics["rmse_fm"][0], np.sqrt(250.0)) self.assertEqual(metrics["mae_fm"][1], 6.0) + def test_profile_tables_derive_consistently_from_families(self) -> None: + for profile in FULL_VALIDATION_PROFILES.values(): + with self.subTest(profile=profile.name): + self.assertEqual( + set(profile.metric_key_map), set(profile.prefactor_by_metric) + ) + self.assertEqual( + set(profile.metric_family_by_key.values()), + set(profile.unit_by_family), + ) + for metric, key in profile.metric_key_map.items(): + self.assertEqual( + profile.metric_family_by_key[key], metric.split(":")[0] + ) + + def test_second_rank_column_follows_the_selected_metric(self) -> None: + for profile in FULL_VALIDATION_PROFILES.values(): + with self.subTest(profile=profile.name): + default = [label for label, _ in profile.columns("e:mae")] + self.assertIn("S_MAE", default) + self.assertNotIn("V_MAE", default) + selected = [label for label, _ in profile.columns("v:rmse")] + self.assertIn("V_RMSE", selected) + self.assertNotIn("S_RMSE", selected) + # Every selectable metric stays reachable by the + # best-checkpoint selector, and one trained quantity never + # occupies two columns. + quantities = {family.prefactors for family in profile.families} + for metric, key in profile.metric_key_map.items(): + columns = profile.columns(metric) + self.assertIn(key, [column_key for _, column_key in columns]) + self.assertEqual(len(columns), 2 * len(quantities)) + + def test_virial_metric_selectors_remain_accepted(self) -> None: + for metric in ("V:MAE", "V:RMSE", "v:mae", "v:rmse"): + self.assertTrue(is_valid_full_validation_metric(metric)) + + def test_energy_profile_reports_stress_and_per_atom_virial(self) -> None: + # A 2 Angstrom cubic cell has volume 8, so a unit virial error becomes + # 1/8 in stress and 1/natoms in per-atom virial. + prediction = { + "energy": np.zeros((1, 1)), + "force": np.zeros((1, 12)), + "virial": np.ones((1, 9)), + } + test_data = { + "find_energy": 1.0, + "find_force": 1.0, + "find_virial": 1.0, + "energy": np.zeros((1, 1)), + "force": np.zeros((1, 12)), + "virial": np.zeros((1, 9)), + "box": np.tile((np.eye(3) * 2.0).reshape(9), (1, 1)), + } + metrics = compute_full_validation_energy_metrics( + prediction, test_data, natoms=4, has_pbc=True + ) + self.assertAlmostEqual(metrics["mae_v_per_atom"][0], 0.25) + self.assertAlmostEqual(metrics["rmse_v_per_atom"][0], 0.25) + self.assertAlmostEqual(metrics["mae_s"][0], 0.125) + self.assertAlmostEqual(metrics["rmse_s"][0], 0.125) + # The two presentations carry the same virial error under different + # normalizations. + self.assertAlmostEqual( + metrics["mae_s"][0] * 8.0, metrics["mae_v_per_atom"][0] * 4.0 + ) + + def test_singular_cell_drops_stress_but_keeps_virial(self) -> None: + prediction = { + "energy": np.zeros((1, 1)), + "force": np.zeros((1, 12)), + "virial": np.ones((1, 9)), + } + test_data = { + "find_energy": 1.0, + "find_force": 1.0, + "find_virial": 1.0, + "energy": np.zeros((1, 1)), + "force": np.zeros((1, 12)), + "virial": np.zeros((1, 9)), + "box": np.zeros((1, 9)), + } + metrics = compute_full_validation_energy_metrics( + prediction, test_data, natoms=4, has_pbc=True + ) + self.assertIn("mae_v_per_atom", metrics) + self.assertNotIn("mae_s", metrics) + def test_spin_profile_omits_magnetic_force_when_unavailable(self) -> None: prediction = { "energy": np.array([[3.0]]), diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py index f52b4f346a..f0ec50b1a8 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -211,6 +211,21 @@ def test_sharing_keeps_the_branch_local_exclusion_mask(self) -> None: assert replica.exclude_types == [[0, 1]] assert replica.readout is base.readout + def test_compression_rejects_excluded_pairs(self) -> None: + """The fused kernel has no type-exclusion branch. + + Compression must refuse rather than emit an artifact that can never + reach the fused path: the re-export would fall back to the plain graph + lower, which the Kokkos spin pair style in turn refuses to load. + """ + excluded = self.build(precision="float32", exclude_types=[[0, 1]]) + with pytest.raises(ValueError, match="type-exclusion branch"): + excluded.enable_compression(0.5) + # The exclusion is the only thing standing in the way. + included = self.build(precision="float32") + included.enable_compression(0.5) + assert included.compress + def test_serialization_preserves_parameters(self) -> None: restored = DescrptDPA4C.deserialize(self.descriptor.serialize()).to(env.DEVICE) original_parameters = dict(self.descriptor.named_parameters()) diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index 856bc0724a..6ad48e936f 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -23,6 +23,11 @@ from deepmd.pt_expt.fitting import ( EnergyFittingNet, ) +from deepmd.pt_expt.infer.charge_state import ( + charge_states, + charge_states_per_frame, + single_charge_state, +) from deepmd.pt_expt.model import ( EnergyModel, ) @@ -55,6 +60,81 @@ def _assert_repeatable(a, b) -> None: np.testing.assert_allclose(a, b, rtol=1e-10, atol=1e-10) +class TestChargeStateBoundary(unittest.TestCase): + """The host boundary must reject a state no declared table row answers. + + A model that gathers table rows declares their ranges, and neither the + gather nor the compiled kernel bounds-checks the index, so a fractional or + out-of-range request has to fail here rather than be truncated into a + neighbouring row or read past the table. A model that embeds the condition + continuously declares no ranges and constrains only the width. + """ + + WIDTH = 2 + + #: Ranges a table-indexing model declares, as DPA4C does. + RANGES = ((-100, 100), (0, 100)) + + #: Requests that address no declared row, with the value each one violates. + UNADDRESSABLE_STATES = ( + ([0.5, 1.0], "first charge_spin value.*must be an integer"), + ([1.0, 2.5], "second charge_spin value.*must be an integer"), + ([float("nan"), 1.0], "first charge_spin value.*must be an integer"), + ([1.0, float("inf")], "second charge_spin value.*must be an integer"), + ([-101.0, 1.0], r"first charge_spin value must lie in \[-100, 100\)"), + ([100.0, 1.0], r"first charge_spin value must lie in \[-100, 100\)"), + ([0.0, -1.0], r"second charge_spin value must lie in \[0, 100\)"), + ([0.0, 100.0], r"second charge_spin value must lie in \[0, 100\)"), + ) + + def test_folded_path_rejects_unaddressable_states(self) -> None: + for state, message in self.UNADDRESSABLE_STATES: + with self.subTest(state=state): + with self.assertRaisesRegex(ValueError, message): + single_charge_state(state, self.WIDTH, self.RANGES) + + def test_input_tensor_path_rejects_unaddressable_states(self) -> None: + for state, message in self.UNADDRESSABLE_STATES: + with self.subTest(state=state): + with self.assertRaisesRegex(ValueError, message): + charge_states_per_frame(state, 1, self.WIDTH, self.RANGES) + + def test_an_unaddressable_state_is_rejected_in_any_frame(self) -> None: + # The check covers every frame, not just the first. + with self.assertRaisesRegex(ValueError, "must be an integer"): + charge_states_per_frame( + [[1.0, 2.0], [0.5, 2.0]], 2, self.WIDTH, self.RANGES + ) + + def test_a_continuous_condition_is_left_alone(self) -> None: + # A model that declares no ranges embeds the condition continuously, + # so a fractional state is legitimate and must survive untouched. + np.testing.assert_array_equal( + charge_states_per_frame([0.5, 0.8], 1, self.WIDTH), + [[0.5, 0.8]], + ) + self.assertEqual(single_charge_state([0.5, 0.8], self.WIDTH), (0.5, 0.8)) + + def test_addressable_states_pass_through_unchanged(self) -> None: + np.testing.assert_array_equal( + charge_states_per_frame([[-3, 4], [0, 1]], 2, self.WIDTH, self.RANGES), + [[-3.0, 4.0], [0.0, 1.0]], + ) + self.assertEqual( + single_charge_state([[2, 1], [2, 1]], self.WIDTH, self.RANGES), (2.0, 1.0) + ) + + def test_shape_contracts_are_kept(self) -> None: + with self.assertRaisesRegex(ValueError, "whole number of 2-wide"): + charge_states([1.0], self.WIDTH) + with self.assertRaisesRegex(ValueError, "one charge state per frame"): + charge_states_per_frame([[1, 2]], 2, self.WIDTH) + with self.assertRaisesRegex(ValueError, "not all equal"): + single_charge_state([[1, 2], [3, 4]], self.WIDTH) + with self.assertRaisesRegex(ValueError, "model indexes 3 tables"): + charge_states([1, 2], self.WIDTH, ((-1, 1), (-1, 1), (-1, 1))) + + class TestDeepEvalEner(unittest.TestCase): """Test pt_expt inference for energy models.""" From a81184fd6b2fcfa0471d53dfc4d2833942efedf1 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 17 Aug 2026 01:04:55 +0800 Subject: [PATCH 10/10] test(pt): compare the relative-force loss on the training device The loss is assembled on the training device while the case builds its inputs on the host, so the assertion compared a CUDA scalar against a host one and failed on any GPU runner. The values already agreed. --- source/tests/pt/test_loss_padding.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 4d0a21ca60..00f50889a3 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -922,7 +922,9 @@ def test_relative_force_norm_uses_normalized_residual(self, masked): actual = _ener_loss_fn(loss_obj, model_pred, label, 2) label_norm = torch.linalg.vector_norm(label_force, dim=-1) residual_norm = label_norm / (label_norm + relative_f) - expected = residual_norm.mean() + # The loss is assembled on the training device, while the inputs above + # are built on the host like every other case in this file. + expected = residual_norm.mean().to(actual.device) torch.testing.assert_close(actual, expected) def test_no_op_for_non_mixed(self):