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 0406eb4b91..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()) @@ -182,6 +188,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..4bd06c764e 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. @@ -731,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/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/__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..62e59cb88b --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4_nn/mlp.py @@ -0,0 +1,247 @@ +# 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_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. + + 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..cdffbef95b --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -0,0 +1,2059 @@ +# 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 ( + CHARGE_STATE_TABLE_RANGES, + 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: + 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. + 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. + 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``. + + 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", + "spin_pair", + "spin_type", + "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, + 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 === + if spin is not None: + 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: + 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.") + 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) + 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 + 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 = ( + 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 === + # 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, + spin_channels=self.spin_channels, + 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), + ) + 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), + ) + ) + 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, + # so the non-scalar block is exactly its complement and needs no + # 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( + 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, + spin: Array | None = None, + charge_spin: Array | 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. + 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. + 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 + ------- + 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. + + Raises + ------ + ValueError + 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 === + 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, + 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: + 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 + + 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, + 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. + + 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 + Frame-level total charge and spin multiplicity with shape + ``(F, 2)``. Read only when the descriptor is configured with + ``add_chg_spin_ebd``. + + 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 + 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(), + None, + self.require_charge_spin(charge_spin, nf, graph.edge_vec), + ) + + # === 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, + spin: Array | None = None, + charge_spin: Array | None = None, + ) -> 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)``. + 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 + ------- + descriptor + 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. + """ + 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 === + # 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] + + # === 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. + conditioned_spin = ( + None + if self.spin is None + else self.spin.conditioned_spin(self.require_spin(spin), atype) + ) + + # === Step 4. Build the masked edge amplitudes and harmonics === + amplitude, basis, envelope, spin_payload = self.build_edge_features( + graph, + atype, + pair_latent, + edge_hidden_bias, + conditioned_spin, + ) + + # === Step 5. 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 6. Build calibrated invariant features === + return ( + self.build_invariant_descriptor( + moments, + center_type_embedding, + divisors, + ), + 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_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. + + 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). + + 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_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. + + 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,)``. + 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, + ) + + # === 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, 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 === + # 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 mixing is not None: + modes = self.radial_mode_head(radial_hidden) # (E, R) + amplitude = amplitude + xp.sum(mixing * modes[:, None, :], axis=-1) + + # === 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, + ) + ) + + # === 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( + self, + amplitude: Array, + basis: Array, + envelope: Array, + spin_payload: Array | None, + spin_onsite: Array | None, + 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,)``. + 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 + Number of output nodes ``N``. + + Returns + ------- + moments + Flat normalized moments with shape ``(N, S)``, where + ``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 + 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 + 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) + 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, + 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) + 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) + 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 + # 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", + "spin", + "charge_spin_embedding", + ): + 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; ``use_spin`` fixes both the presence + 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 + 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), + None if self.use_spin is None else tuple(self.use_spin), + self.add_chg_spin_ebd, + 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. 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, + ) + + del path + sampled = merged() if callable(merged) else merged + 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 + 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) + # 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: + 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 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) + ) + 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, + box, + self.get_rcut(), + ) + output, _ = self.call_graph( + 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, + self.get_dim_out(), + ) + if output_np.shape[0] == 0: + continue + 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] + ) + finally: + self.mean, self.stddev = mean_backup, stddev_backup + + if not np.any(active_count > 0.0): + return + # 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 + finite_positive = np.isfinite(feature_rms) & (feature_rms > 0.0) + invalid = measured & ( + ~np.isfinite(feature_rms) | (geometric & ~finite_positive) + ) + if np.any(invalid): + raise ValueError( + "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 <= 0.0: + raise ValueError( + f"DPA4C type embedding has a degenerate calibration RMS {target_rms}" + ) + # 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) + + # 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_count = active_count[mass] + mass_mean = value_sum[mass] / mass_count + mass_stddev = np.sqrt( + np.maximum(square_sum[mass] / mass_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) + 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. + + 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``, + 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 ``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 + ------ + 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)) + ) + # 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, + 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], + "charge_spin": ( + None if charge_spin is None else charge_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: + """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, + "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(), + "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") + spin_channels = data.pop("spin_channels") + charge_spin_embedding = data.pop("charge_spin_embedding") + + 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.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"], + ) + 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``. + + 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. + """ + 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.""" + 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. + + 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 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 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 + + 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..a3610297f4 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py @@ -0,0 +1,56 @@ +# 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 .charge_state import ( + CHARGE_STATE_TABLE_RANGES, + ChargeStateEmbedding, + canonicalize_charge_spin, + validate_charge_state, +) +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, +) +from .spin import ( + NEIGHBOR_QUADRUPOLE_CHANNELS, + SpinChannels, + derive_spin_channels, +) + +__all__ = [ + "CHARGE_STATE_TABLE_RANGES", + "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/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/charge_state.py b/deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py new file mode 100644 index 0000000000..d0d8a0bf54 --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/charge_state.py @@ -0,0 +1,383 @@ +# 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) + +#: 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. + + 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_STATE_FIELDS, + CHARGE_STATE_TABLE_RANGES, + strict=True, + ): + 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( + 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/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..c9a28ae23b --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/pair_film.py @@ -0,0 +1,374 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Ordered type-pair FiLM cache for DPA4C.""" + +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.common import ( + to_numpy_array, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +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},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}),\\ + \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 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`, :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 + ---------- + 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. + spin_channels + Number :math:`C_s` of ordered spin scale and shift channels. Zero + omits the spin tables. + precision + Parameter precision. + trainable + Whether the pair-encoder weights are trainable. + seed + Random seed. + + Raises + ------ + ValueError + 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, + 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}") + 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) + 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=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) + + @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 + ---------- + type_embedding + Complete type table with shape ``(T + 1, channels)``, where the + trailing row is the zero padding type. + + Returns + ------- + 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) + 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), + ) + 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 + + 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, + ) + ) + + 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: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 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. + + Returns + ------- + dict[str, Any] + 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 + 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") + 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/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/descriptor/dpa4c_nn/spin.py b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py new file mode 100644 index 0000000000..56ea73389f --- /dev/null +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/spin.py @@ -0,0 +1,802 @@ +# 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. + """ + + CONFIG_DERIVED_ARRAYS = ("spin_mask",) + + 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 (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, + ) + # 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) + + # === 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 + # 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, + ) -> 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 + Per-edge ordered spin scales with shape ``(E, spin_channels)``. + pair_shift + Per-edge ordered spin shifts with the same shape. + + 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) + 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), + 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 + ), + "spin_gate": to_numpy_array(self.spin_gate), + }, + } + + @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 = dict(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 + ) + 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. + + 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..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. @@ -344,6 +355,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/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 8459d0847f..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,10 +501,9 @@ def call( ) # [nf, nloc, 3] huber_ncomp = 3 else: - diff_3 = xp.reshape(force_hat - force, (*_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, ) @@ -529,19 +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 - force, (-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 @@ -554,9 +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 - force, (*_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. @@ -571,7 +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 - force, (-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/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 b6f765c22a..9300eafcb2 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, @@ -1159,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 6c02b1566e..2fd9b84b02 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"), @@ -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, @@ -513,18 +517,19 @@ 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 - # convention). + # Stress sigma = -virial / volume, in eV/ų (tensile-positive + # 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/__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..997fd14215 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( @@ -454,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_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 27a8c958c8..fd9e734464 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( @@ -1036,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, @@ -1046,6 +1022,7 @@ def dpa1_graph_compress_energy_force( graph.source_order, graph.source_row_ptr, graph.n_node, + 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 a257a96c92..40b5e536f4 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( @@ -269,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, @@ -279,6 +276,7 @@ def _cpu( source_order, source_row_ptr, n_node, + edge_vec.to(fprec).new_empty(0), node_capacity, do_atomic_virial, ) @@ -371,8 +369,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 +424,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..ceae538ca7 --- /dev/null +++ b/deepmd/kernels/cuda/dpa4c/canonical.py @@ -0,0 +1,489 @@ +# 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 the complete compact DPA4C operator suite is 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, + ) + 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, energy_gradient) + ) + + +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, + spin: torch.Tensor, + spin_pair: torch.Tensor, + spin_type: 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, + spin_pair, + spin_type, + 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), spin.ndim == 2 + ) + 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, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del descriptor_gradient, state + return _backward_shapes(edge_vec, args[13]) + + +def _backward_inplace_fake( + descriptor_gradient: torch.Tensor, + state: torch.Tensor, + edge_vec: torch.Tensor, + *args: Any, +) -> 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.ndim != 2: + 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( + edge_vec: torch.Tensor, + source: torch.Tensor, + destination_row_ptr: torch.Tensor, + atype: torch.Tensor, + *args: Any, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + del source, destination_row_ptr + spin = args[10] + 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 ( + 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), + ) + + +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 + + 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, spin_gradient, edge_spin_gradient = _cpu_backward( + gradient, state, *descriptor_args + ) + return energy, edge_gradient, spin_gradient, edge_spin_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 = 23 + + +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, + 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 + 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. + spin + Per-node magnetic moments with shape ``(N, 3)`` for a spin-conditioned + descriptor, or ``None``. + + 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. + 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 + ------ + 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, + 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) + # 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, + ) + 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 new file mode 100644 index 0000000000..ff1a918ed4 --- /dev/null +++ b/deepmd/kernels/cuda/dpa4c/graph_compress.py @@ -0,0 +1,2130 @@ +# 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 ( + TYPE_CHECKING, + 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, + 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", + "descriptor_profile", + "dpa4c_graph_compress", + "dpa4c_graph_compress_energy_force", + "ef_op_available", + "ensure_registered", + "fitting_energy_and_gradient", + "mega_eligible", + "op_available", + "reduce_edge_spin_gradient", +] + +SUPPORTED_CHANNELS = (8, 16, 32, 64, 128) +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. +_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. + + 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 + 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. + spin_width + Width of the trailing native-spin invariant block. + """ + + channels: int + lmax: int + degree_channels: tuple[int, ...] + ranks: tuple[int, ...] + moment_width: int + output_width: int + gram_base: int + bispectrum_base: int + spin_channels: int + spin_width: 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 + + @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] + return slice(stop - self.spin_width, stop) + + +def descriptor_profile( + channels: int, + lmax: int, + has_spin: bool = False, +) -> DescriptorProfile: + """Derive every compiled width from the 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 + ---------- + channels + 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 + ------- + 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 + 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] + + spin_dim + + 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, + spin_channels=spin_channels, + spin_width=spin_dim, + ) + + +@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 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( + "DPA4C compressed CUDA requires descriptor precision `float32`, " + f"got {sample_parameter.dtype}" + ) + if not mega_eligible(descriptor): + raise ValueError( + "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, + descriptor.spin is not None, + ) + table, info = build_radial_table(descriptor, stride) + + with torch.no_grad(): + 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) + ) + 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, + profile.lmax, + ) + to_device = {"device": device} + return { + "data": table, + "info": info, + "spin_type": _build_spin_type_cache(descriptor, device), + "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(), + **build_charge_state_artifacts(descriptor, descriptor.default_chg_spin), + } + + +def charge_state_artifacts( + descriptor: Any, + 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] + + 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. 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. + device + Device that receives the packed table. + + Returns + ------- + torch.Tensor + Per-type table with shape ``(T + 1, 4)``, or an empty tensor for a + spin-free descriptor. + """ + if descriptor.spin is None: + 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) + 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() + ) + + +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, + *, + 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. + + 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] + 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]) + + # === 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. 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( + [ + blocks[0][:, 0, :], + *[_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], + ], + dim=-1, + ) + 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, + 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[: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].ndim == 2, + ) + 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, + spin: torch.Tensor, + spin_pair: torch.Tensor, + spin_type: 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, + spin_pair, + spin_type, + canonical, + table_stride, + table_max, + rcut, + eps, + degree_floor, + ) + profile = descriptor_profile( + int(type_embedding.shape[1]), int(lmax), spin.ndim == 2 + ) + 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, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del descriptor_gradient, state + spin = args[15] + 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 ( + 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( + descriptor_gradient: torch.Tensor, + state: torch.Tensor, + edge_vec: torch.Tensor, + *args: Any, +) -> 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 + spin = args[15] + 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(): + descriptor = _cpu_descriptor( + value, + *args[:15], + *args[18:], + spin=moment, + spin_pair=args[16], + spin_type=args[17], + ) + 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[:19]) + ctx.scalars = inputs[19:] + 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: + """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, + )[0] + return (edge_gradient,) + (None,) * 25 + + +_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, + 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 native spin block. + """ + empty = descriptor.compress_spin_type.new_empty(0) + 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, + 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. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor. + graph + NeighborGraph with destination CSR topology. + atype + Flat node types with shape ``(N,)``. + spin + Per-node magnetic moments with shape ``(N, 3)``, or ``None``. + + 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") + 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, 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) + + +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, + 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 + ---------- + 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. + spin + Per-node magnetic moments with shape ``(N, 3)``, or ``None``. + + 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. + force_mag + Per-node magnetic force with shape ``(N, 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, spin), + 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, + spin_gradient, + edge_spin_gradient, + ) = torch.ops.deepmd.dpa4c_graph_compress_backward( + descriptor_gradient, + state, + edge_vec, + *operator_args, + ) + # 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, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + graph.n_node, + edge_spin_gradient, + node_capacity, + do_atomic_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( + 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..9e8e2ecd93 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,72 @@ 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 _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, @@ -60,14 +128,18 @@ 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, 3) + if _has_spin_cotangent(edge_spin_gradient) + else g_e.new_empty(0), ) @@ -78,15 +150,19 @@ 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, 3) + if _has_spin_cotangent(edge_spin_gradient) + else g_e.new_empty(0), ) @@ -100,9 +176,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( @@ -115,7 +192,18 @@ def _cpu( ) if not want_atom_virial: atom_virial = atom_virial.new_zeros(0, 3, 3) - return force, atom_virial, virial + 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 + 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_empty(0) + return force, atom_virial, virial, magnetic_force def _canonical_cpu( @@ -125,9 +213,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( @@ -164,13 +253,14 @@ 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, source_order, source_row_ptr, n_node_per_frame, + edge_spin_gradient, node_capacity, want_atom_virial, ) @@ -192,8 +282,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", @@ -212,9 +306,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: @@ -239,6 +334,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 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 @@ -253,6 +351,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 a + rank-one empty sentinel when no spin cotangent was supplied. """ ensure_registered() return torch.ops.deepmd.edge_force_virial( @@ -265,6 +366,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, ) @@ -277,9 +379,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 @@ -294,6 +397,9 @@ 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 a rank-one + empty sentinel when the model carries no magnetic degree of freedom. node_capacity Flat node count ``N``. want_atom_virial @@ -301,8 +407,10 @@ def canonical_edge_force_virial( Returns ------- - tuple[torch.Tensor, torch.Tensor, torch.Tensor] - Force, optional atom virial, and frame virial. + 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 a rank-one empty sentinel when no + spin cotangent was supplied. """ ensure_registered() return torch.ops.deepmd.canonical_edge_force_virial( @@ -312,6 +420,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/kernels/cuda/graph_fitting.py b/deepmd/kernels/cuda/graph_fitting.py index fb8b7624c5..c22ea4ad06 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,27 +54,58 @@ ) __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.""" - 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: """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 +134,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 +148,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 +240,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 +254,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 +285,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 +352,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 +413,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 +506,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/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/__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/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index 91266b8788..c3f3632413 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)) + + @BaseDescriptor.register("se_atten") @BaseDescriptor.register("dpa1") @torch_module @@ -1068,15 +1091,17 @@ 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. 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`. @@ -1096,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, @@ -1123,7 +1152,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 +1184,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 new file mode 100644 index 0000000000..b399a17eaa --- /dev/null +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -0,0 +1,642 @@ +# 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",), + "OrderedPairFiLM": ( + "adam_spin_scale_anchor", + "adam_spin_shift_anchor", + ), + "SpinChannels": ( + "adam_spin_vector_weight", + "adam_spin_quadrupole_weight", + "spin_gate", + ), +} + + +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, + spin: torch.Tensor | None = None, + charge_spin: torch.Tensor | 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. + 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 + ------- + 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): + # 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() + return super().call_graph( + graph, + atype, + type_embedding=type_embedding, + comm_dict=comm_dict, + spin=spin, + charge_spin=charge_spin, + ) + + def build_edge_features( + self, + graph: Any, + *args: Any, + ) -> 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 + 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. 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 + 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, the ordered pair cache, and the conditioned spin, + 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,)``. + spin_payload + Masked per-edge spin payload, or ``None``. + """ + 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( + 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. + + ``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, and every + spin-conditioned descriptor, keeps the default behavior exactly. + """ + 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: + 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 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) + 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 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, + 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. + + 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 + 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, 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: + 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, + spin: torch.Tensor | None = None, + ) -> ( + tuple[ + torch.Tensor, + 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. The + trailing output is the magnetic force, empty for a spin-free model. + """ + 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, + 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/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/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 b084f3fdef..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, ) @@ -122,19 +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 _warn_legacy_edge_vec(metadata: dict) -> None: """Warn once per model load when an edge_vec-schema artifact is opened. @@ -223,6 +215,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) @@ -247,7 +242,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 +311,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: " @@ -493,7 +496,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 @@ -536,6 +541,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 ( @@ -696,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", } @@ -914,25 +924,139 @@ 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 _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. + + 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, self._chg_spin_table_ranges()) + ) + 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 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, ) @@ -945,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() @@ -1753,7 +1883,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 ) @@ -1811,7 +1945,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 @@ -2090,22 +2224,73 @@ 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." + ) + # 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, + ) + 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: @@ -2204,7 +2389,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,11 +2404,14 @@ 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)." ) + # 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, ) @@ -2566,6 +2757,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 +2830,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/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 7838781889..791ce3da42 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_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 6aa5f4add4..28853ecd11 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, @@ -66,8 +136,9 @@ 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 DPA1 deployment graph. + """Evaluate an eligible compressed canonical deployment graph. Parameters ---------- @@ -78,7 +149,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 @@ -90,6 +161,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 ------- @@ -97,17 +171,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 +209,35 @@ 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, force_mag = ( + dpa4c_canonical_compress_energy_force( + descriptor, + fitting, + graph, + atype, + 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, + 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, @@ -148,6 +245,15 @@ def forward_lower_canonical_graph( "virial": virial, "mask": output_mask.to(torch.int32), } + 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 return result @@ -236,6 +342,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 @@ -254,45 +374,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, @@ -304,6 +390,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. @@ -330,6 +417,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 ------- @@ -346,8 +435,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, @@ -425,8 +519,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 +794,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/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 7aa82215b2..bbd3262342 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.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, + # 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..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, @@ -1496,6 +1497,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 +1618,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 +2374,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..e1a593ccbc 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 = " " @@ -279,13 +279,13 @@ 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 ) 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() @@ -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 } @@ -524,11 +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 = ( - not self.profile.needs_spin - and 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/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..2269c47b7c 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 the compact canonical trace inputs for compressed DPA1.""" + """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,17 +677,29 @@ 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, + ) + 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) - return ( + edge_storage_dim = torch.export.Dim( + "nedge_storage", + min=2, + max=UINT32_MAX, + ) + shapes = ( {0: node_dim}, {0: nframes_dim}, {0: nframes_dim}, @@ -691,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( @@ -922,8 +941,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 @@ -1006,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. @@ -1021,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: @@ -1061,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 @@ -1069,6 +1127,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 @@ -1123,6 +1186,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 +1315,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 +1333,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" @@ -1320,23 +1393,28 @@ 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', 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 # 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: @@ -1412,17 +1490,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; @@ -1458,12 +1537,12 @@ def _trace_and_export( # 2. Collect metadata metadata = _collect_metadata( model, - is_spin=is_spin, + spin_scheme=spin_scheme, lower_kind=lower_kind, ) # 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 +1588,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 +1599,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 @@ -1575,6 +1659,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, @@ -1962,6 +2047,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, @@ -2012,6 +2280,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 @@ -2064,11 +2344,22 @@ 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_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 # ``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: @@ -2121,3 +2412,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/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 d9d57dfb1a..f19240ab45 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 [ @@ -451,6 +452,149 @@ def descrpt_se_a_args() -> list[Argument]: ] +@descrpt_args_plugin.register( + "dpa4c", + alias=["DPA4C"], + doc=supported_backends("pt_expt") + + "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( + "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], + optional=True, + default=None, + doc="Random seed for parameter initialization.", + ), + ] + + @descrpt_args_plugin.register( "dpa4", alias=["DPA4", "SeZM", "sezm"], @@ -5930,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 d75c0b7385..efcd4de8c4 100644 --- a/deepmd/utils/eval_metrics.py +++ b/deepmd/utils/eval_metrics.py @@ -22,6 +22,14 @@ "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"), @@ -77,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, @@ -143,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], @@ -155,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( @@ -174,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( @@ -190,6 +240,7 @@ def compute_energy_type_metrics( force=force, virial=virial, virial_per_atom=virial_per_atom, + stress=stress, ) @@ -282,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 ------- @@ -301,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``. Spin models do not report virial, so no virial - metric is produced. + 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 ---------- @@ -316,26 +366,17 @@ 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 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( @@ -344,43 +385,56 @@ def compute_full_validation_spin_metrics( 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 @@ -391,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[ @@ -403,97 +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"), - ("V_MAE", "mae_v_per_atom"), - ("V_RMSE", "rmse_v_per_atom"), + 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", - "v:mae": "mae_v_per_atom", - "v:rmse": "rmse_v_per_atom", - }, - 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", - }, - unit_by_family={ - "e": ("meV/atom", 1000.0), - "f": ("meV/Å", 1000.0), - "v": ("meV/atom", 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"), - }, 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" + "# 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"), + 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", - }, - 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", - }, - unit_by_family={ - "e": ("meV/atom", 1000.0), - "fr": ("meV/Å", 1000.0), - "fm": ("meV/μB", 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"), - }, 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 " + _SECOND_RANK_NOTE ), compute_system_metrics=compute_full_validation_spin_metrics, ) diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md new file mode 100644 index 0000000000..e75f286b54 --- /dev/null +++ b/doc/model/dpa4c.md @@ -0,0 +1,467 @@ +# 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. + +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. + +## 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. + +## How it works + +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. + +Three consequences shape how the model is used in practice. + +- **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 + +### 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. + +### Options that matter + +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] +> 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 + +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 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. + +Train, freeze, then compress the frozen archive: + +```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 +``` + +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. + +Compression requires: + +- the PyTorch Exportable backend on CUDA; +- `precision: "float32"`; +- `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. + +`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 +``` + +### 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 + +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} . +``` + +### 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 +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 | + +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. + +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 +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 a spin model in LAMMPS + +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 +pair_style dpa4spin compressed_model.pt2 +pair_coeff * * Ni O +``` + +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. + +A worked example, a rocksalt NiO cell in its type-II antiferromagnetic order, +is provided in `examples/spin/dpa4c/lmp/`. + +## 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. + +## 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 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 + +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 +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`. + +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. 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/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..aa4153f4f6 --- /dev/null +++ b/examples/spin/dpa4c/lmp/README.md @@ -0,0 +1,69 @@ +# 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` 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 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 +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..050ce53903 --- /dev/null +++ b/examples/spin/dpa4c/lmp/in.lammps @@ -0,0 +1,62 @@ +# 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 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: +# +# 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/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..b658821082 --- /dev/null +++ b/examples/water/dpa4c/input.json @@ -0,0 +1,91 @@ +{ + "_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..d98d744387 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. @@ -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 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( @@ -450,11 +452,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); @@ -1255,6 +1257,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 30 + */ +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 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 30 + */ +extern bool DP_DeepSpinUsesNativeSpinScheme(DP_DeepSpin* dp); + /** * @brief Evaluate the energy, force and virial by using a DP with the mixed *type. (double version) @@ -2559,6 +2614,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. @@ -2614,6 +2683,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. @@ -2712,6 +2795,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. @@ -2773,6 +2870,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/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/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 08cabbcf0c..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. * @@ -1351,11 +1364,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) { @@ -1889,6 +1902,64 @@ 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. + * + * 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. @@ -2537,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. @@ -3020,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 4e91591ad3..c35a4a1763 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -773,6 +773,39 @@ bool validate_model_devi_nframes(DP_DeepBaseModelDevi* dp, const int nframes) { } // namespace +namespace deepmd { +namespace c_api_internal { + +/** + * @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)); + } + 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 c_api_internal +} // namespace deepmd + template void DP_DeepPotModelDeviCompute_variant( DP_DeepPotModelDevi* dp, @@ -1789,11 +1822,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) { @@ -2173,6 +2206,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, @@ -2717,8 +2790,24 @@ 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(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(); } +void DP_DeepSpinSetChargeSpin(DP_DeepSpin* dp, + const double* charge_spin, + const int numb_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) { return DP_DeepBaseModelIsAParamNAll(static_cast(dp)); } @@ -2759,6 +2848,14 @@ 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(deepmd::c_api_internal::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)); @@ -2838,6 +2935,14 @@ 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(deepmd::c_api_internal::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_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/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/DeepPot.h b/source/api_cc/include/DeepPot.h index 262c653713..56cbef8704 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" @@ -205,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 @@ -357,11 +379,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,17 +851,23 @@ 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); 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; }; @@ -884,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 a4ea314003..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); @@ -348,11 +376,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,19 +433,28 @@ 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); 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 2a13b79cf8..bf57849b14 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" @@ -167,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 @@ -252,6 +274,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; }; /** @@ -538,12 +596,72 @@ 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. **/ 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; }; @@ -588,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 new file mode 100644 index 0000000000..51b0c2dd69 --- /dev/null +++ b/source/api_cc/include/NativeSpinPTExpt.h @@ -0,0 +1,497 @@ +// 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" + +// 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; +} + +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 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 + * 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(); + ~NativeSpinPTExpt() override; + 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 = "") override; + + double cutoff() const override { + assert(inited); + return rcut; + }; + int numb_types() const override { + assert(inited); + return ntypes; + }; + int numb_types_spin() const override { + assert(inited); + return ntypes_spin; + }; + int dim_fparam() const override { + assert(inited); + return dfparam; + }; + int dim_aparam() const override { + 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 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) override; + bool is_aparam_nall() const override { return false; }; + bool has_default_fparam() const override { + 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) 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 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 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 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 + // 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. + * + * 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, 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& charge_spin); + + /** + * @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. + * + * 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. + * @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; + // 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; + // 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; + // 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 + +#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 066c6ba991..a9894e2f04 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -3,9 +3,11 @@ #ifdef BUILD_PYTORCH #include +#include #include #include +#include #include #include #include @@ -195,6 +197,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; @@ -234,8 +248,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 { @@ -393,27 +418,52 @@ 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 = 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"); + } + const auto index_type = canonicalGraphIndexType(); auto source = torch::zeros({storage_count}, - graph.edge_index.options().dtype(torch::kInt64)); + 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)); + auto source_order = + torch::arange(storage_count, + graph.edge_index.options().dtype(torch::kInt64)) + .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::kInt64)); + .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::kInt64)); + .copy_(graph.source_order.slice(0, 0, edge_count).to(index_type)); } return {graph.atype, graph.n_node, @@ -843,6 +893,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/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 4888e0b387..3373d6fe71 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) { @@ -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 b2ad13336c..a2d3eeddcd 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -9,10 +9,12 @@ #include #include #include +#include #include #include #include #include +#include #include "SimulationRegion.h" #include "common.h" @@ -22,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; @@ -189,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(); @@ -238,7 +245,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 +274,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(); @@ -356,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) { @@ -376,6 +415,31 @@ 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)); + } + // 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( const torch::Tensor& coord, const torch::Tensor& atype, @@ -659,6 +723,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); @@ -717,6 +785,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 @@ -740,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 " @@ -758,6 +838,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 @@ -783,11 +878,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/ @@ -815,7 +911,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_) { @@ -825,7 +921,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 { @@ -865,37 +961,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 @@ -1215,6 +1297,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 @@ -1377,6 +1462,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. @@ -1517,37 +1606,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 @@ -1608,6 +1683,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); } @@ -1682,21 +1760,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(); @@ -1730,9 +1798,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, @@ -1896,21 +1964,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(); @@ -1946,9 +2004,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, @@ -2472,11 +2530,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 +2546,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 +2560,13 @@ 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(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), - {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 +2575,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 +2665,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/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..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(); @@ -465,6 +469,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/DeepSpinPTExpt.cc b/source/api_cc/src/DeepSpinPTExpt.cc index 958b240ead..dcf2075b19 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, } } + // 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 " + + 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 new file mode 100644 index 0000000000..adf308052e --- /dev/null +++ b/source/api_cc/src/NativeSpinPTExpt.cc @@ -0,0 +1,1248 @@ +// 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::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; + +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 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. + * + * 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); +} + +/** + * @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) { + 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; + 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(); + 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 || dchgspin > 0) { + throw deepmd::deepmd_exception( + "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."); + } + } + // 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); + } + } + + // 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) { + 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::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)); + } + // 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) { + 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, + const torch::Tensor& charge_spin) { + 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); + } + if (dchgspin > 0) { + inputs.push_back(charge_spin); + } + 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), + make_chg_spin_tensor(default_chg_spin_, dchgspin, 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; + 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. + 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_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, " + "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; + 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 + // 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::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, + 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(deepmd::canonicalGraphIndexType()) + .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/api_cc/src/commonPTExpt.h b/source/api_cc/src/commonPTExpt.h index dfd6057c59..bb67ffd0f5 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,133 @@ 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"); + } + // 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()) { + constants[constants_[ii]] = tables[ii]; + } + } + target.load_constants(constants, /*use_inactive=*/true, + /*check_full_update=*/true); + runner->swap_constant_buffer(); + runner->free_inactive_constant_buffer(); + } + + 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/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc index ef83eae28b..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::kInt64); + 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::kInt64))); + 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 new file mode 100644 index 0000000000..38389d6bd5 --- /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 = 0; + 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.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 9ecd68f055..fe02503576 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 @@ -30,10 +29,8 @@ using namespace LAMMPS_NS; 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), @@ -204,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); @@ -225,107 +218,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_owner.extent(0) < nall) { - k_owner = DAT::tdual_int_1d("deepmd/kk:owner", 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; - 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; - k_owner.template modify(); - k_owner.template sync(); - d_owner = k_owner.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 = @@ -335,15 +238,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; - const bool multi = multi_rank; - auto owner = d_owner; - auto loc2model = d_loc2model; - 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( @@ -368,7 +283,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 +366,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; } @@ -475,7 +390,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); @@ -493,186 +408,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 owner = d_owner; - auto loc2model = d_loc2model; - const bool multi = multi_rank; - 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::int64_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) { - 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) { - ++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); - const std::size_t required = static_cast(storage_count); - if (workspace.edge_capacity < required) { - const std::size_t slack = required / 8 + 64; - if (required > std::numeric_limits::max() - slack) { - error->one(FLERR, "Compact DPA1 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; - Kokkos::parallel_for( - "deepmd/kk:canonical_fill", 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 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; - } - } - }); - - 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 = 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; - } - }); - 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; - 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; - }); - 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) = 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; - }); - } - return edge_count; -} - template void PairDeepMDKokkos::compute(int eflag, int vflag) { if (!device_path_ok) { @@ -696,14 +431,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 @@ -759,8 +494,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) { @@ -780,12 +515,13 @@ 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; } - 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 @@ -796,29 +532,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()); @@ -835,7 +572,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; @@ -923,7 +660,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 d09e9a7abc..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,18 +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; - 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 @@ -80,40 +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 - // 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). - // 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) - 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 @@ -124,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_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 new file mode 100644 index 0000000000..0a8d6e1ba8 --- /dev/null +++ b/source/lmp/pair_dpa4spin.cpp @@ -0,0 +1,583 @@ +// 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 first 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; + + // 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); + } +} + +/* ---------------------------------------------------------------------- + map the atom types onto the elements of the model +------------------------------------------------------------------------- */ + +void PairDPA4Spin::coeff(int narg, char** arg) { + if (narg < 2) { + error->all(FLERR, "Incorrect args for pair coefficients"); + } + 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. 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] += dforce[3 * ii + dd] * force_scale; + fm[ii][dd] += dforce_mag[3 * ii + dd] * sp[ii][3] * magnetic_force_scale; + } + } + + 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..7d0aba8977 --- /dev/null +++ b/source/lmp/pair_dpa4spin.h @@ -0,0 +1,109 @@ +// 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 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. +class PairDPA4Spin : public Pair { + public: + PairDPA4Spin(class LAMMPS*); + ~PairDPA4Spin() override; + + // 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; + // 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/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 5e1a46ec91..b1d6c31cc3 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..0c57805a47 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,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_idts, 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()); @@ -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, @@ -131,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,9 +151,9 @@ 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..45eb90dc69 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress.cu @@ -0,0 +1,956 @@ +// 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; + int spin_channels; +}; + +template +Dimensions dimensions_of() { + using P = Profile; + return {P::MomentWidth, P::OutputWidth, P::C1, + deepmd_dpa4c::coupling_record_count(Lmax), P::Cs}; +} + +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, 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 + 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; + // 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; + 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.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( + 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"); + 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 || + 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: 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.dim() == 2; + 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.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; + 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); + 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, + 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) { + 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 + // 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), 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()); + 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}; +} + +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.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. + 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), 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()); + 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"); + 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) { + 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 = + 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); + // 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()), spin_gradient, + edge_spin_gradient}; +} + +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, spin, + spin_pair, spin_type, 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, + 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) { + 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, spin, spin_pair, spin_type, true, lmax, table_stride, + table_max, rcut, eps, degree_floor); +} + +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"); + 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, spin, + spin_pair, spin_type, true, lmax, table_stride, table_max, rcut, eps, + degree_floor, reuse_state); +} + +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, spin, spin_pair, spin_type, lmax, table_stride, table_max, + rcut, eps, degree_floor, false); +} + +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, 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 +// 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, + 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, + 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 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); + 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); + // 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()), spin_gradient, + edge_spin_gradient}; + } + 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, 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(); + 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(); + 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()), spin_gradient, + edge_spin_gradient}; +} + +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, " + "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)"); + 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, 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( + "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, " + "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, + &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, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, " + "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_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, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, " + "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_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, " + "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 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 new file mode 100644 index 0000000000..002fb5bd29 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress.cuh @@ -0,0 +1,737 @@ +// 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; + // Retained for the native spin branch, which gathers the source moment. + long source; +}; + +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 = source; + 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]); +} + +// === 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 { + 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) { + // 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); +} + +// 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) { + // Geometric widths only: the probe layout has no spin block. + 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. +// +// 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 (&direction_gradient)[3], + float radial_gradient, + float (&output)[3]) { + const float dux = + 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 = + 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 = + 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; + 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..c2f94f33f9 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress_kernel.cuh @@ -0,0 +1,2132 @@ +// 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; + +// === 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, + 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; + 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; + 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] = {}; + + // 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 === + 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; + + // 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; + 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 (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 + // 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]); + } + } + 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]; + { + 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; + } + } + } + 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(); + + 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]); + } + + // === 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 === + +// 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) { + // 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; + 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); + } + 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; + } + } + 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]); + } + } + } + + // === 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 === + // 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, + 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; + 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); + } + // 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] = {}; +#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); + } + } + } + + // 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]) { + 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; + 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; + } + 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; + 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; + } + } + 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); + } + + // 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) { + 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 (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; + } + } + } + + 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, + direction_gradient); + } + } + } + + // 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, direction_gradient, 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]; + } + 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]; + } + } + } +} + +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, + 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; + 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; + 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; + } + } +} + +// === 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 { + // 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; + 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 + // 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, + HasSpin ? args.edge_spin_gradient : nullptr); + } + } +}; + +// 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); + } + } +} + +// 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> +void dispatch_modes(const Arguments& args, cudaStream_t stream) { + if (args.radial_modes > 0) { + dispatch_spin(args, stream); + } else { + dispatch_spin(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..70e52f1bd6 --- /dev/null +++ b/source/op/pt/dpa4c_graph_compress_launch.h @@ -0,0 +1,398 @@ +// 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; + // 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; + 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; + + // === 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 === + +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. +// +// 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. +// 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; + +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; +}; + +// ``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. +// ``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); + static constexpr int C2 = degree_two_width(Channels); + static constexpr int K1 = C2; + static constexpr int K2 = 2; + + // === 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; + 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 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. + 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); + // 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 + SpinDim; + 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 = + HasSpin ? EdgeMap::SpinForward : EdgeMap::Forward; + static constexpr int BackwardEdgeWidth = + HasSpin ? EdgeMap::SpinBackward : EdgeMap::Backward; + static constexpr int NodeWidth = + NodeLanes != 0 ? NodeLanes : kNodeLanesNarrow; + 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. +#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..844d46eb56 100644 --- a/source/op/pt/edge_force_virial.cu +++ b/source/op/pt/edge_force_virial.cu @@ -18,7 +18,9 @@ #include #include +#include #include +#include namespace { @@ -26,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(); \ @@ -49,7 +63,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, @@ -59,8 +73,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; @@ -72,6 +88,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) { @@ -107,6 +130,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]); @@ -127,6 +155,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] += @@ -143,30 +176,41 @@ __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; + } } } -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 +227,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 +249,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, @@ -221,50 +312,54 @@ 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, + bool has_spin, 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()); + 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(), + has_spin ? edge_spin_gradient.data_ptr() : nullptr, + force.data_ptr(), node_virial.data_ptr(), + has_spin ? 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"); - 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( - 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); @@ -274,19 +369,15 @@ std::tuple assemble_force_virial( ? atom_virial : torch::empty({node_count, 3, 3}, options); auto virial = torch::zeros({frame_count, 3, 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}; + return {force, atom_virial, virial, magnetic_force}; } - 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,19 +386,28 @@ 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, - source_row_ptr, frame_row_ptr, force, node_virial, virial_partial, - virial, stream); + 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, 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, has_spin, + force, node_virial, magnetic_force, 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, - source_row_ptr, frame_row_ptr, force, node_virial, virial_partial, - virial, stream); + 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, has_spin, + force, node_virial, magnetic_force, virial_partial, virial, + stream); } }); - return {force, atom_virial, virial}; + return {force, atom_virial, virial, magnetic_force}; } } // namespace @@ -358,18 +458,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() && @@ -393,23 +494,29 @@ 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"); + 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, - 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(); @@ -430,20 +537,69 @@ 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 " "entries"); + TORCH_CHECK( + 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()); 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 +// 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) { @@ -457,22 +613,30 @@ 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( + "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..b4564784db 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__ float2 act_vg(float z) { +__device__ __forceinline__ float act_value(float z) { if constexpr (ACT == 0) { - float a = tanhf(z); - return make_float2(a, 1.f - a * a); + return tanhf(z); } - 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))); + return z * sigmoid(z); } -// 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). template -__global__ void layer_epilogue_kernel(long total4, +__device__ __forceinline__ float act_derivative(float z) { + if constexpr (ACT == 0) { + const float a = tanhf(z); + return 1.f - a * a; + } + 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) (+ 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 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,274 @@ __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; +} + +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 +// 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) { + 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); - 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]; 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 +475,125 @@ 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) { + 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); + 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 && + 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; + 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..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 @@ -96,38 +97,91 @@ 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, -// 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 rank-two ``edge_spin_gradient`` adds the per-source total +// of the magnetic cotangent, which shares the source grouping the force +// 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, + 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 new file mode 100644 index 0000000000..9832dabd8c --- /dev/null +++ b/source/tests/common/dpmodel/test_descriptor_dpa4c.py @@ -0,0 +1,1800 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import dataclasses +import math +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.descriptor.dpa4c_nn.spin import ( + SpinChannels, +) +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.pair_latent(descriptor.type_embedding.call()), + )[:3] + + +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, + None, + None, + 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 + + +# === 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. + + 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, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 23, + "use_spin": [True, False], + } + config.update(overrides) + descriptor = DescrptDPA4C(**config) + if descriptor.spin is not None: + descriptor.spin.spin_gate[...] = 1.0 + return descriptor + + +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 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() + 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.pair_latent( + self.descriptor.type_embedding.call() + ), + None, + 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) + + +#: 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/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/common/test_examples.py b/source/tests/common/test_examples.py index 111436fd3b..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", @@ -65,6 +66,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/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_loss_padding.py b/source/tests/pt/test_loss_padding.py index 927c9ba0cd..00f50889a3 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,34 @@ 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) + # 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): """All-ones mask gives same force loss as no mask.""" f = _t(NP, 3) 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_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py index e14dc81610..e3bdef7ab8 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, @@ -1273,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, @@ -1283,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_empty(0), n, True, ) @@ -1432,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) @@ -1447,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, @@ -1457,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_empty(0), n_node, True, ) @@ -1504,7 +1539,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, @@ -1592,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, @@ -1602,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_empty(0), n, True, ) @@ -1705,6 +1741,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, @@ -1924,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, @@ -1934,9 +1971,11 @@ def _fused( src_order, src_row_ptr, n_node, + edge_vec.new_empty(0), total, True, ) + return force, atom_virial, virial def _assert_device_parity(self, device) -> None: args = self._random_graph(torch.device(device)) @@ -2016,6 +2055,7 @@ def test_compact_canonical_parity(self) -> None: source_order, source_row_ptr, n_node, + edge_vec.new_empty(0), total, True, ) @@ -2026,12 +2066,110 @@ def test_compact_canonical_parity(self) -> None: compact.source_row_ptr, compact.source_order, compact.n_node, + edge_vec.new_empty(0), 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_empty(0), + total, + True, + ) + self.assertEqual(tuple(empty.shape), (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) + + # 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.py b/source/tests/pt_expt/descriptor/test_dpa4c.py new file mode 100644 index 0000000000..f0ec50b1a8 --- /dev/null +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -0,0 +1,656 @@ +# 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_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()) + 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.pair_latent(descriptor.type_embedding.call()), + )[:3] + 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, + ) + + +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.""" + + 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() + # 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( + 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 + 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 + + 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 new file mode 100644 index 0000000000..83f1d13417 --- /dev/null +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py @@ -0,0 +1,1548 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Numerical contract of the compressed DPA4C CUDA mega kernel.""" + +import dataclasses +from collections.abc import ( + Sequence, +) + +import numpy as np +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, + _cpu_forward, + _table_lookup, + build_compression_artifacts, + build_radial_table, + descriptor_profile, + dpa4c_graph_compress_energy_force, + ensure_registered, + mega_eligible, + 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"], + 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]) +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, *_spin_free(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) + # 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) + 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, *_spin_free(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) + + +def _build_charge_descriptor( + channels: int = 8, + radial_modes: int = 0, + default_chg_spin: Sequence[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=( + None if default_chg_spin is None else list(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) + 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, *_spin_free(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", "shape", "dtype", "message"), + [ + (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, + 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] = torch.zeros(shape, dtype=dtype, device="cuda") + 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:18], + *arguments[19:], + ) + 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:18], + *arguments[19:], + ) + 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, + edge_vec.new_empty(0), + 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) + + +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. + + 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. 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( + rcut=3.0, + ntypes=2, + channels=channels, + lmax=lmax, + n_radial=8, + radial_modes=radial_modes, + precision="float32", + seed=17, + use_spin=[True, False], + ) + .to(device) + .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 + + +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"), + [(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/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.""" 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") 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 new file mode 100644 index 0000000000..68394c32f1 --- /dev/null +++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py @@ -0,0 +1,713 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from types import ( + SimpleNamespace, +) + +import numpy as np +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, +) + +_GPU = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="CUDA is required for compact canonical graph export", +) + + +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" + # 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, + dtype: torch.dtype = torch.float64, +) -> dict[str, torch.Tensor]: + sample = build_synthetic_graph_inputs( + model, + e_max=None, + nframes=2, + nloc=7, + dtype=dtype, + 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_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) + 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" + + +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, + 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_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, + ) + + 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) + # 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, dtype=torch.float32) + monkeypatch.setenv("DP_CUDA_INFER", "2") + actual = _run_graph(model, dtype=torch.float32) + for key in ( + "energy", + "energy_redu", + "energy_derv_r", + "energy_derv_c", + "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" +) +@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. + + 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. + + 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) + + 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"] 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", + ), } 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)