diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index ee406c11ae..70f752de62 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -320,6 +320,51 @@ def forward_common_atomic( atom_mask = xp_take_first_n(ext_atom_mask, 1, nloc) return self._finalize_atomic_ret(ret_dict, atom_mask, atype) + def _prepare_graph_nodes( + self, + n_node: Array, + n_local: Array | None, + atype: Array, + reference: Array, + ) -> tuple[Array, Array]: + """Apply node masks shared by generic and compact graph forwards.""" + xp = array_api_compat.array_namespace(reference) + atype = xp.asarray(atype, device=array_api_compat.device(reference)) + atom_mask = self.make_atom_mask(atype) + atype_clamped = xp.where(atom_mask, atype, xp.zeros_like(atype)) + output_mask = atom_mask + if n_local is not None: + from deepmd.dpmodel.utils.neighbor_graph import ( + node_ownership_mask, + ) + + output_mask = output_mask & node_ownership_mask( + n_node, + n_local, + atype.shape[0], + ) + if self.atom_excl is not None: + output_mask = xp.logical_and( + output_mask, + self.atom_excl.build_type_exclude_mask(atype_clamped), + ) + return atype_clamped, output_mask + + def _prepare_graph_inputs( + self, + graph: "NeighborGraph", + atype: Array, + ) -> tuple["NeighborGraph", Array, Array]: + """Apply graph masks shared by standard and fused atomic forwards.""" + atype_clamped, output_mask = self._prepare_graph_nodes( + graph.n_node, + graph.n_local, + atype, + graph.edge_vec, + ) + self._assert_graph_pair_excluded(graph, atype_clamped) + return graph, atype_clamped, output_mask + def forward_common_atomic_graph( self, graph: "NeighborGraph", @@ -359,17 +404,7 @@ def forward_common_atomic_graph( the result dict on the flat node axis, defined by the `FittingOutputDef`. """ - xp = array_api_compat.array_namespace(graph.edge_vec) - atype = xp.asarray(atype, device=array_api_compat.device(graph.edge_vec)) - atom_mask = self.make_atom_mask(atype) # (N,) bool - atype_clamped = xp.where(atom_mask, atype, xp.zeros_like(atype)) - # NOTE: model-level ``pair_exclude_types`` is NOT applied here. It is a - # graph-BUILD transform (decision #18) already folded into - # ``graph.edge_mask`` by the NeighborGraph builder (Python) or - # ``applyPairExclusion`` (C++); this method consumes a pre-excluded graph. - # Fail-safe (eager only): guard against a caller that skipped the build - # seam, which would silently INCLUDE excluded pairs (fail-open). - self._assert_graph_pair_excluded(graph, atype_clamped) + graph, atype_clamped, output_mask = self._prepare_graph_inputs(graph, atype) ret_dict = self.forward_atomic_graph( graph, atype_clamped, @@ -377,7 +412,7 @@ def forward_common_atomic_graph( aparam=aparam, charge_spin=charge_spin, ) - return self._finalize_atomic_ret(ret_dict, atom_mask, atype) + return self._finalize_atomic_ret(ret_dict, output_mask, atype) def _assert_nlist_pair_excluded(self, nlist: Array, extended_atype: Array) -> None: """Fail-safe: assert the nlist reaching the dense seam is pre-excluded. @@ -491,10 +526,16 @@ def _finalize_atomic_ret( """ xp = array_api_compat.array_namespace(atype) - ret_dict = self.apply_out_stat(ret_dict, atype) + safe_atype = xp.where( + self.make_atom_mask(atype), + atype, + xp.zeros_like(atype), + ) + ret_dict = self.apply_out_stat(ret_dict, safe_atype) if self.atom_excl is not None: atom_mask = xp.logical_and( - atom_mask, self.atom_excl.build_type_exclude_mask(atype) + atom_mask, + self.atom_excl.build_type_exclude_mask(safe_atype), ) lead = atom_mask.shape # (nf, nloc) dense | (N,) graph for kk in ret_dict.keys(): diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 440eb75284..ba8b224d19 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -37,6 +37,46 @@ ) +def _extend_graph_aparam( + aparam: Array, + n_node: Array, + n_local: Array, + n_total: int, +) -> Array: + """Expand frame-local atomic parameters onto a local-plus-halo node axis.""" + import array_api_compat + + from deepmd.dpmodel.utils.neighbor_graph import ( + frame_id_from_n_node, + node_ownership_mask, + ) + + xp = array_api_compat.array_namespace(aparam, n_node, n_local) + frame_id = frame_id_from_n_node(n_node, n_total=n_total) + frame_end = xp.cumulative_sum(n_node) + frame_start = frame_end - n_node + node_index = xp.arange( + n_total, + dtype=n_node.dtype, + device=array_api_compat.device(n_node), + ) + index_in_frame = node_index - xp.take(frame_start, frame_id, axis=0) + local_capacity = aparam.shape[1] + sentinel = xp.zeros( + (aparam.shape[0], 1, aparam.shape[2]), + dtype=aparam.dtype, + device=array_api_compat.device(aparam), + ) + padded_aparam = xp.concat([aparam, sentinel], axis=1) + padded_capacity = local_capacity + 1 + local_index = index_in_frame % padded_capacity + flat_index = frame_id * padded_capacity + local_index + flat_aparam = xp.reshape(padded_aparam, (-1, aparam.shape[-1])) + gathered = xp.take(flat_aparam, flat_index, axis=0) + ownership = node_ownership_mask(n_node, n_local, n_total) + return xp.where(ownership[:, None], gathered, xp.zeros_like(gathered)) + + @BaseAtomicModel.register("standard") class DPAtomicModel(BaseAtomicModel): r"""Model give atomic prediction of some physical property. @@ -302,10 +342,27 @@ def forward_atomic_graph( ) fparam_node = None if fparam is not None: - frame_id = frame_id_from_n_node(graph.n_node) + frame_id = frame_id_from_n_node( + graph.n_node, + n_total=atype.shape[0], + ) fparam_node = xp.take(fparam, frame_id, axis=0) # (N, ndf) + aparam_node = aparam + if aparam is not None and graph.n_local is not None and aparam.ndim == 3: + aparam_node = _extend_graph_aparam( + aparam, + graph.n_node, + graph.n_local, + atype.shape[0], + ) return self.fitting_net.call_graph( - gg, atype, gr=rot_mat, g2=None, h2=None, fparam=fparam_node, aparam=aparam + gg, + atype, + gr=rot_mat, + g2=None, + h2=None, + fparam=fparam_node, + aparam=aparam_node, ) def compute_or_load_stat( diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 3a47a7319c..b338a88044 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -434,12 +434,13 @@ def uses_graph_lower(self) -> bool: """Returns whether this descriptor supports the graph-native lower. The graph-native lower (``call_graph``) covers the factorizable path - AND transformer attention (``attn_layer >= 0``, NeighborGraph PR-D) - with concat OR strip type-embedding. ``exclude_types`` is fully - supported via + and transformer attention (``attn_layer >= 0``) with concat or strip + type embedding. ``exclude_types`` is applied through :func:`~deepmd.dpmodel.utils.neighbor_graph.apply_pair_exclusion`. - Compressed descriptors are the remaining ineligible config and fall - back to the legacy dense path, so those models keep working unchanged. + Geo-compressed strip models + (``geo_compress=True``, ``attn_layer == 0``) are also graph-eligible; + tebd-only compression and compressed descriptors with descriptor-level + exclusions stay on the legacy dense path. Eligibility does NOT imply numerical interchangeability with the dense route for every config: with ``smooth_type_embedding=True`` @@ -449,13 +450,13 @@ def uses_graph_lower(self) -> bool: """ if self._graph_lower_disabled: return False - # compressed descriptors have no graph kernel (geo/tebd tabulation is - # dense-only); keep them on the legacy dense path. if self.compress: - return False - # strip is graph-eligible (per-edge factorized embedding, no neighbor - # coupling); exclude_types is graph-native via ``apply_pair_exclusion`` - # (owned at this seam). Only compression / the disable flag force dense. + return ( + self.geo_compress + and self.se_atten.tebd_input_mode == "strip" + and self.se_atten.attn_layer == 0 + and not self.se_atten.exclude_types + ) return self.se_atten.tebd_input_mode in ("concat", "strip") def disable_graph_lower(self) -> None: @@ -1775,7 +1776,7 @@ def call_graph( Notes ----- Known limitations: - - ``tebd_input_mode`` in {"concat", "strip"}; compressed descriptors stay dense; + - ``tebd_input_mode`` in {"concat", "strip"}; - ``exclude_types`` is applied graph-natively via ``apply_pair_exclusion``. """ from deepmd.dpmodel.utils.neighbor_graph import ( diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 3a5d19f83a..3f4465a2de 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -666,9 +666,9 @@ def forward_common_atomic_graph( Parameters ---------- atype - (N,) flat LOCAL atom types, ``N == sum(n_node)``. + (N,) flat local-plus-halo atom types, ``N == sum(n_node)``. n_node - (nf,) per-frame local atom counts. + (nf,) per-frame total node counts. edge_index (2, E) ``[src, dst]`` edge endpoints (flat local indices). edge_vec @@ -676,18 +676,15 @@ def forward_common_atomic_graph( edge_mask (E,) boolean/0-1 valid-edge mask. n_local - Per-rank local atom counts for multi-rank inference. Ignored in - PR-A (single-rank); accepted for ABI stability. + Per-frame owned node counts. Halo fitting outputs are masked. fparam Frame parameter, ``(nf, ndf)``. aparam Atomic parameter, ``(N, nda)``. comm_dict - MPI communication metadata. Ignored in PR-A; accepted for ABI - stability. + Optional MPI communication metadata. charge_spin - charge/spin conditioning. Ignored in PR-A; accepted for ABI - stability with charge/spin-conditioned descriptors. + Charge/spin conditioning. Returns ------- @@ -701,6 +698,7 @@ def forward_common_atomic_graph( edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, + n_local=n_local, ) atomic_ret = self.atomic_model.forward_common_atomic_graph( graph, atype, fparam=fparam, aparam=aparam, charge_spin=charge_spin @@ -739,9 +737,9 @@ def call_common_lower_graph( Parameters ---------- atype - (N,) flat LOCAL atom types, ``N == sum(n_node)``. + (N,) flat local-plus-halo atom types, ``N == sum(n_node)``. n_node - (nf,) per-frame local atom counts. + (nf,) per-frame total node counts. edge_index (2, E) ``[src, dst]`` edge endpoints (flat local indices). edge_vec @@ -749,18 +747,15 @@ def call_common_lower_graph( edge_mask (E,) boolean/0-1 valid-edge mask. n_local - Per-rank local atom counts for multi-rank inference. Ignored in - PR-A (single-rank); accepted for ABI stability. + Per-frame owned node counts. Halo fitting outputs are masked. fparam Frame parameter, ``(nf, ndf)``. aparam Atomic parameter, ``(N, nda)``. comm_dict - MPI communication metadata. Ignored in PR-A; accepted for ABI - stability. + Optional MPI communication metadata. charge_spin - charge/spin conditioning. Ignored in PR-A; accepted for ABI - stability with charge/spin-conditioned descriptors. + Charge/spin conditioning. Returns ------- diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 4e53cbe789..c047d37e97 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -3,8 +3,9 @@ The unified edge/graph neighbor-list contract and its supporting machinery: ``graph`` (the ``NeighborGraph``/``GraphLayout`` contract + derived node-validity -+ edge padding), ``builder`` (the carry-all ``build_neighbor_graph`` dispatcher + -the ``from_dense_quartet`` legacy converter), ``segment`` (mask-aware ++ edge padding), ``csr`` (backend-agnostic CSR construction and canonicalization), +``builder`` (the carry-all ``build_neighbor_graph`` dispatcher + the +``from_dense_quartet`` legacy converter), ``segment`` (mask-aware segment-reduction toolkit), and ``derivatives`` (edge force/virial assembly). See the design discussion wanghan-iapcm/deepmd-kit#4. """ @@ -25,6 +26,11 @@ from_dense_quartet, graph_from_dense_quartet, ) +from .csr import ( + attach_edge_csr, + build_edge_csr, + canonicalize_neighbor_graph, +) from .derivatives import ( edge_force_virial, ) @@ -39,6 +45,7 @@ NeighborGraph, apply_pair_exclusion, frame_id_from_n_node, + node_ownership_mask, node_validity_mask, pad_and_guard_angles, pad_and_guard_edges, @@ -61,9 +68,12 @@ "angle_to_node_sum", "apply_pair_exclusion", "attach_angles", + "attach_edge_csr", "build_angle_index", + "build_edge_csr", "build_neighbor_graph", "build_neighbor_graph_ase", + "canonicalize_neighbor_graph", "center_edge_pairs", "edge_env_mat", "edge_force_virial", @@ -72,6 +82,7 @@ "graph_angle_cos", "graph_from_dense_quartet", "neighbor_graph_from_ijs", + "node_ownership_mask", "node_validity_mask", "pad_and_guard_angles", "pad_and_guard_edges", diff --git a/deepmd/dpmodel/utils/neighbor_graph/angles.py b/deepmd/dpmodel/utils/neighbor_graph/angles.py index afe2677c86..d6414c19ff 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/angles.py +++ b/deepmd/dpmodel/utils/neighbor_graph/angles.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """3-body angle graph: pairs of edges sharing a center within a_rcut. -Angles reference EDGES (angle_index into [0,E)); edge_vec stays the only -geometry leaf. a_sel is normalization-only (not a truncation). Reuses PR-D's -center_edge_pairs; a_rcut filters the participating edges. +Angles reference edges (angle_index into [0,E)); edge_vec stays the only +geometry leaf. a_sel is normalization-only (not a truncation), and a_rcut +filters the edges passed to center_edge_pairs. """ from __future__ import ( @@ -89,8 +89,8 @@ def build_angle_index( # the normalization below. dist = safe_for_vector_norm(edge_vec, axis=-1) # (E,) a_edge_mask = xp.astype(edge_mask, xp.bool) & (dist < a_rcut) - # compact eager form only (static_nnei not exposed until angle export is - # needed, PR-G). dst = edge_index[1, :] per the [src, dst] SoA convention. + # The compact eager form groups by destination under the [src, dst] SoA + # convention. q_e, k_e, pair_mask = center_edge_pairs( edge_index[1, :], a_edge_mask, diff --git a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py index fa163634a4..8c7288c7b4 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py @@ -21,6 +21,9 @@ import numpy as np +from .csr import ( + attach_edge_csr, +) from .from_ijs import ( neighbor_graph_from_ijs, ) @@ -47,6 +50,8 @@ def build_neighbor_graph_ase( rcut: float, layout: GraphLayout | None = None, *, + with_csr: bool = False, + canonicalize: bool = False, pair_excl: PairExcludeMask | None = None, compact: bool = False, ) -> NeighborGraph: @@ -73,6 +78,12 @@ def build_neighbor_graph_ase( cutoff radius. layout edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. + with_csr + Whether to construct destination/source CSR views for a consumer that + requires edge-grouped reductions. + canonicalize + Whether to reorder every edge field into destination-major form. Implies + ``with_csr=True``. pair_excl Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` for model-level ``pair_exclude_types``. When given, @@ -152,7 +163,14 @@ def _to_cpu_numpy(x: Any) -> np.ndarray: S_all, nframe_all = S_all[keep], nframe_all[keep] graph = neighbor_graph_from_ijs( - i_all, j_all, S_all, coord, box, nframe_all, nloc, layout=layout + i_all, + j_all, + S_all, + coord, + box, + nframe_all, + nloc, + layout=layout, ) if pair_excl is not None: import array_api_compat @@ -160,4 +178,6 @@ def _to_cpu_numpy(x: Any) -> np.ndarray: xp = array_api_compat.array_namespace(coord) atype_flat = xp.reshape(xp.asarray(atype), (-1,)) graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) + if with_csr or canonicalize: + graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize) return graph diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index 54f70e0bc8..f948714e4a 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -43,6 +43,10 @@ xp_take_first_n, ) +from .csr import ( + attach_edge_csr, + build_edge_csr, +) from .graph import ( GraphLayout, NeighborGraph, @@ -66,6 +70,9 @@ def from_dense_quartet( mapping: Array, layout: GraphLayout | None = None, compact: bool = True, + *, + with_csr: bool = False, + canonicalize: bool = False, ) -> NeighborGraph: """Convert a legacy extended quartet into a ghost-free NeighborGraph (CONVERTER). @@ -107,6 +114,12 @@ def from_dense_quartet( ``src`` pointing at the center (in-range, masked) -- so no ``nonzero`` is used and the converter is jit/export-traceable. The masked edges contribute zero in a downstream ``segment_sum``, so the descriptor output is unchanged. + with_csr + Whether to construct destination/source CSR views for a consumer that + requires edge-grouped reductions. + canonicalize + Whether to reorder every edge field into destination-major form. Implies + ``with_csr=True``. Returns ------- @@ -117,6 +130,7 @@ def from_dense_quartet( """ if layout is None: layout = GraphLayout() + with_csr = with_csr or canonicalize xp = array_api_compat.array_namespace(extended_coord, nlist, mapping) dev = array_api_compat.device(extended_coord) nf, nloc, nsel = nlist.shape @@ -157,11 +171,38 @@ def from_dense_quartet( edge_index = xp.astype(xp.stack([src, dst], axis=0), xp.int64) edge_mask = valid n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) + if not with_csr: + return NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + ) + ( + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + edge_index, + edge_vec, + edge_mask, + nf * nloc, + canonicalize=canonicalize, + ) return NeighborGraph( n_node=n_node, edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=canonicalize, ) else: # COMPACT: drop invalid slots via nonzero (dynamic shape -> eager only, @@ -198,11 +239,38 @@ def from_dense_quartet( edge_index, edge_vec, layout.edge_capacity, layout.min_edges ) n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) + if not with_csr: + return NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + ) + ( + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + edge_index, + edge_vec, + edge_mask, + nf * nloc, + canonicalize=canonicalize, + ) return NeighborGraph( n_node=n_node, edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=canonicalize, ) @@ -265,6 +333,8 @@ def build_neighbor_graph( rcut: float, layout: GraphLayout | None = None, *, + with_csr: bool = False, + canonicalize: bool = False, pair_excl: PairExcludeMask | None = None, compact: bool = False, ) -> NeighborGraph: @@ -305,6 +375,12 @@ def build_neighbor_graph( at non-binding ``sel``). layout edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. + with_csr + Whether to construct destination/source CSR views for a consumer that + requires edge-grouped reductions. + canonicalize + Whether to reorder every edge field into destination-major form. Implies + ``with_csr=True``. pair_excl Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` for model-level ``pair_exclude_types``. When given, @@ -323,6 +399,7 @@ def build_neighbor_graph( if layout is None: layout = GraphLayout() + with_csr = with_csr or canonicalize xp = array_api_compat.array_namespace(coord, atype) dev = array_api_compat.device(coord) nf, nloc = atype.shape[:2] @@ -384,4 +461,6 @@ def build_neighbor_graph( if pair_excl is not None: atype_flat = xp.reshape(atype, (-1,)) graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) + if with_csr: + graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize) return graph diff --git a/deepmd/dpmodel/utils/neighbor_graph/csr.py b/deepmd/dpmodel/utils/neighbor_graph/csr.py new file mode 100644 index 0000000000..d499a8f369 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/csr.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-agnostic CSR topology helpers for :class:`NeighborGraph`.""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + replace, +) +from typing import ( + TYPE_CHECKING, +) + +import array_api_compat + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + + from .graph import ( + NeighborGraph, + ) + + +def build_edge_csr( + edge_index: Array, + edge_vec: Array, + edge_mask: Array, + n_nodes: int, + canonicalize: bool = False, +) -> tuple[Array, Array, Array, Array, Array, Array, Array]: + """Build destination/source CSR views of an edge payload. + + By default, both views store permutations into the original edge stream. + With ``canonicalize=True``, a stable destination permutation is applied to + every edge field, masked entries move to the suffix, and + ``destination_order`` becomes the identity. Stable ordering preserves the + incoming order within each destination segment. + + Parameters + ---------- + edge_index : Array + Edge endpoints with shape ``(2, E)`` in ``[source, destination]`` order. + edge_vec : Array + Edge vectors with shape ``(E, 3)``. + edge_mask : Array + Real-edge mask with shape ``(E,)``. + n_nodes : int + Number of nodes in the flat graph. + canonicalize : bool, default: False + Whether to reorder the payload into destination-major form. + + Returns + ------- + edge_index : Array + Edge endpoints with shape ``(2, E)``. + edge_vec : Array + Edge vectors with shape ``(E, 3)``. + edge_mask : Array + Real-edge mask. + destination_order : Array + Edge indices grouped by destination with shape ``(E,)`` and the same + dtype as ``edge_index``. + destination_row_ptr : Array + Destination CSR offsets with shape ``(N + 1,)`` and dtype int64. + source_order : Array + Edge indices grouped by source with shape ``(E,)`` and the same dtype + as ``edge_index``. + source_row_ptr : Array + Source CSR offsets with shape ``(N + 1,)`` and dtype int64. + + Raises + ------ + OverflowError + If an int32 edge payload has more than ``2**31 - 1`` entries. + """ + xp = array_api_compat.array_namespace(edge_index) + device = array_api_compat.device(edge_index) + edge_count = edge_index.shape[1] + if ( + edge_index.dtype == xp.int32 + and isinstance(edge_count, int) + and edge_count > 2**31 - 1 + ): + raise OverflowError( + "int32 edge payload cannot represent a permutation of more than " + "2**31 - 1 edges" + ) + padding_node = xp.asarray(n_nodes, dtype=edge_index.dtype, device=device) + + destination_key = xp.where(edge_mask, edge_index[1], padding_node) + destination_order = xp.argsort(destination_key, stable=True) + ordered_destination = xp.take(destination_key, destination_order, axis=0) + if canonicalize: + edge_index = xp.take(edge_index, destination_order, axis=1) + edge_vec = xp.take(edge_vec, destination_order, axis=0) + edge_mask = xp.take(edge_mask, destination_order, axis=0) + destination_order = xp.arange( + edge_index.shape[1], dtype=edge_index.dtype, device=device + ) + else: + destination_order = xp.astype(destination_order, edge_index.dtype) + node_boundaries = xp.arange( + n_nodes + 1, + dtype=edge_index.dtype, + device=device, + ) + destination_row_ptr = xp.astype( + xp.searchsorted(ordered_destination, node_boundaries, side="left"), + xp.int64, + ) + + source_key = xp.where(edge_mask, edge_index[0], padding_node) + source_order = xp.argsort(source_key, stable=True) + ordered_source = xp.take(source_key, source_order, axis=0) + source_order = xp.astype(source_order, edge_index.dtype) + source_row_ptr = xp.astype( + xp.searchsorted(ordered_source, node_boundaries, side="left"), + xp.int64, + ) + return ( + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) + + +def attach_edge_csr( + graph: NeighborGraph, + n_nodes: int, + canonicalize: bool = False, +) -> NeighborGraph: + """Attach destination/source CSR views to an edge graph. + + Parameters + ---------- + graph : NeighborGraph + The graph whose current edge payload and mask define the CSR views. + n_nodes : int + Number of nodes on the flat graph axis. + canonicalize : bool, default: False + Whether to reorder the payload into destination-major form. + + Returns + ------- + NeighborGraph + A copy carrying CSR views consistent with its edge payload and mask. + + Raises + ------ + ValueError + If canonicalization is requested for a graph with angle indices. + """ + if canonicalize and graph.angle_index is not None: + raise ValueError("cannot canonicalize a graph with angle indices") + ( + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + n_nodes, + canonicalize=canonicalize, + ) + return replace( + graph, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=canonicalize, + ) + + +def canonicalize_neighbor_graph( + graph: NeighborGraph, + n_nodes: int, +) -> NeighborGraph: + """Return a destination-major edge graph. + + Generic graph builders preserve the incoming edge order. Deployment + adapters call this function when their ABI guarantees a destination-major + payload and identity destination permutation. + + Parameters + ---------- + graph : NeighborGraph + The graph to canonicalize. + n_nodes : int + Number of nodes on the flat graph axis. + + Returns + ------- + NeighborGraph + A graph with every edge field reordered consistently and + ``destination_sorted=True``. + + Raises + ------ + ValueError + If the graph contains angle indices, which would require a matching + edge-index remapping. + """ + return attach_edge_csr(graph, n_nodes, canonicalize=True) diff --git a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py index 599877d7e2..09111fa5d5 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -23,6 +23,9 @@ import array_api_compat +from .graph import ( + frame_id_from_n_node, +) from .segment import ( segment_sum, ) @@ -80,40 +83,22 @@ def edge_force_virial( frame via the frame of their ``dst`` node. """ xp = array_api_compat.array_namespace(g_e) - # node-axis size; when a ``node_capacity`` is supplied (the jax/export path) - # use it AS-IS so we never call int() on the traced ``sum(n_node)`` -- and, - # crucially, never on ``node_capacity`` itself: under symbolic make_fx / - # torch.export it is a SymInt (``atype.shape[0]``); ``int(SymInt)`` would - # SPECIALIZE the node axis to the trace-time sample size, baking a constant - # ``N`` into the scatter and breaking dynamic-``N`` inference. + # A supplied capacity remains symbolic under export; converting it to int + # would specialize the dynamic node axis to the trace sample. n_out = node_capacity if node_capacity is not None else int(xp.sum(n_node)) nf = n_node.shape[0] - # zero padding/guard contributions; cast mask to g's dtype (array-API pure, - # CLAUDE.md mask-multiply guideline — avoids bool*float under array_api_strict) + if isinstance(n_out, int) and n_out == 0: + device = array_api_compat.device(g_e) + return ( + xp.zeros((0, 3), dtype=g_e.dtype, device=device), + xp.zeros((0, 3, 3), dtype=g_e.dtype, device=device), + xp.zeros((nf, 3, 3), dtype=g_e.dtype, device=device), + ) + # Padding edges carry no force or virial contribution. g = g_e * xp.astype(edge_mask[:, None], g_e.dtype) - # Wrap node indices into ``[0, n_out)`` so every scatter address is provably - # in-bounds. For a well-formed graph every real edge already has - # ``index < n_out`` (== ``atype.shape[0]``), so this modulo is the IDENTITY on - # real edges (pinned by test_modulo_clamp_leaves_real_edges_unchanged) -- a - # correctness-preserving guard, not a value fixup. - # - # Why it is needed (root cause, GPU-confirmed): under the dynamic-edge graph - # ``torch.export`` path the node count is traced as several equal-but-distinct - # symbols (``atype.shape[0]``, ``fit_ret.shape[0]``, ...), tied only by - # ``aten._assert_scalar(Eq(...))`` nodes. ``_strip_shape_assertions`` - # (pt_expt/utils/serialization.py) neutralises ALL such asserts so export can - # trace -- which also drops those node-count equalities, so inductor can no - # longer prove the scatter index and its bound ``ks0 == n_out`` share a symbol - # and emits ``tl.device_assert(idx < ks0)`` (fatal on CUDA; unchecked on CPU, - # which is why all CPU dev/CI was green). ``% n_out`` discharges that guard - # unconditionally. This is the PERMANENT fix: the upstream alternative -- - # making the SHARED, spin-export-critical ``_strip_shape_assertions`` - # selective -- risks re-triggering the torch.export bugs it exists to bypass - # and the spin ``.pt2`` path, so it is deliberately NOT taken. - # - # Pure arithmetic => torch.export-safe, unlike ``xp.clip`` (SymInt bound - # breaks array_api_compat's clip) and unlike a mask-multiply (which misses the - # ``edge_mask == 1`` indices the stripped guard mis-bounds). + # Real endpoints are already in range. Modulo leaves them unchanged while + # bounding masked sentinel endpoints after export removes symbolic shape + # equalities between the endpoint and output node axes. src = edge_index[0] % n_out dst = edge_index[1] % n_out # force (output sized to the node axis, incl. any padding tail) @@ -122,14 +107,14 @@ def edge_force_virial( w_edge = -(g[:, :, None] * edge_vec[:, None, :]) # (E, 3, 3) # atom virial: full-to-src atom_virial = segment_sum(w_edge, src, n_out) # (N, 3, 3) - # per-frame virial: assign each edge to the frame of its dst node. Node - # ``k`` belongs to frame ``searchsorted(cumsum(n_node), k, "right")`` because - # real nodes are compact frame-major (frame f owns a contiguous block). - boundaries = xp.cumulative_sum(n_node) # (nf,) per-frame node upper bounds - edge_frame = xp.astype( - xp.searchsorted(boundaries, dst, side="right"), xp.int64 - ) # (E,) in [0, nf] - # wrap into [0, nf) for the same CUDA-bounds reason (export-safe modulo) - edge_frame = edge_frame % nf - virial = segment_sum(w_edge, edge_frame, nf) # (nf, 3, 3) + # Per-frame virial: reduce the PER-ATOM virial by its node frame (N -> nf) + # rather than the per-edge virial by its edge frame (E -> nf). Real edges + # never cross frames (``src`` and ``dst`` share a frame), so summing the + # full-to-``src`` atom virial over a frame's nodes equals summing ``w_edge`` + # over that frame's edges. Reducing the N nodes instead of the E >> N edges + # avoids serializing ~E scatter-adds into a single per-frame accumulator -- + # the single-frame (inference) worst case, where every edge targets one slot + # and the scatter degenerates into a fully contended atomic reduction. + node_frame = frame_id_from_n_node(n_node, n_total=n_out) # (N,) frame per node + virial = segment_sum(atom_virial, node_frame, nf) # (nf, 3, 3) return force, atom_virial, virial diff --git a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py index 0136dd4bf4..d8c8f73c27 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py +++ b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py @@ -24,6 +24,9 @@ import array_api_compat +from .csr import ( + build_edge_csr, +) from .graph import ( GraphLayout, NeighborGraph, @@ -45,6 +48,9 @@ def neighbor_graph_from_ijs( nframe_id: Array, nloc: int, layout: GraphLayout | None = None, + *, + with_csr: bool = False, + canonicalize: bool = False, ) -> NeighborGraph: """Convert a sparse ``(i, j, S)`` edge list into a :class:`NeighborGraph`. @@ -70,6 +76,12 @@ def neighbor_graph_from_ijs( number of local atoms per frame (used for the frame-major node offset). layout edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. + with_csr + Whether to construct destination/source CSR views for a consumer that + requires edge-grouped reductions. + canonicalize + Whether to reorder every edge field into destination-major form. Implies + ``with_csr=True``. Returns ------- @@ -80,6 +92,7 @@ def neighbor_graph_from_ijs( """ if layout is None: layout = GraphLayout() + with_csr = with_csr or canonicalize xp = array_api_compat.array_namespace(coord) dev = array_api_compat.device(coord) nf = coord.shape[0] @@ -108,9 +121,36 @@ def neighbor_graph_from_ijs( edge_index, edge_vec, layout.edge_capacity, layout.min_edges ) n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) + if not with_csr: + return NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + ) + ( + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + edge_index, + edge_vec, + edge_mask, + nf * nloc, + canonicalize=canonicalize, + ) return NeighborGraph( n_node=n_node, edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=canonicalize, ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index 5c05cc8c64..f702c2c6a1 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -14,6 +14,7 @@ from dataclasses import ( dataclass, + field, ) from typing import ( TYPE_CHECKING, @@ -36,7 +37,10 @@ class NeighborGraph: Geometry enters the model ONLY through ``edge_vec`` (the single autograd leaf). ``edge_index``/``angle_index`` use the SoA ``(2, .)`` layout so the - src/dst index vectors are contiguous. + src/dst index vectors are contiguous. Destination/source CSR views address + the current payload through permutations. Builders may preserve incoming + order or apply a stable destination-major canonicalization. Consumers must + apply ``edge_mask`` even inside a CSR row. """ n_node: Array @@ -53,6 +57,16 @@ class NeighborGraph: """(2, A) int [edge_a, edge_b] sharing a center; into [0, E). None if no angles.""" angle_mask: Array | None = None """(A,) bool real vs padding on the angle axis. None if no angles.""" + destination_order: Array | None = field(default=None, kw_only=True) + """(E,) edge permutation grouped by destination; same dtype as edge_index.""" + destination_row_ptr: Array | None = field(default=None, kw_only=True) + """(N + 1,) int64 offsets into ``destination_order``.""" + source_order: Array | None = field(default=None, kw_only=True) + """(E,) source-grouped edge permutation; same dtype as edge_index.""" + source_row_ptr: Array | None = field(default=None, kw_only=True) + """(N + 1,) int64 CSR offsets into ``source_order``.""" + destination_sorted: bool = field(default=False, kw_only=True) + """Whether the payload is destination-major and destination_order is identity.""" @dataclass @@ -221,6 +235,37 @@ def frame_id_from_n_node(n_node: Array, n_total: int | None = None) -> Array: return xp.minimum(frame_id, xp.astype(last_frame, xp.int64)) +def node_ownership_mask(n_node: Array, n_local: Array, n_total: int) -> Array: + """Return the owned-node mask for a local-plus-halo graph. + + Each frame occupies one contiguous block of ``n_node[f]`` nodes, with its + ``n_local[f]`` owned nodes first and halo nodes after them. + + Parameters + ---------- + n_node + Total node counts per frame with shape ``(nf,)``. + n_local + Owned node counts per frame with shape ``(nf,)``. + n_total + Size of the flat node axis. + + Returns + ------- + Array + Boolean ownership mask with shape ``(n_total,)``. + """ + xp = array_api_compat.array_namespace(n_node, n_local) + device = array_api_compat.device(n_node) + node_index = xp.arange(n_total, dtype=n_node.dtype, device=device) + frame_id = frame_id_from_n_node(n_node, n_total=n_total) + frame_end = xp.cumulative_sum(n_node) + frame_start = frame_end - n_node + index_in_frame = node_index - xp.take(frame_start, frame_id, axis=0) + local_count = xp.take(n_local, frame_id, axis=0) + return index_in_frame < local_count + + def apply_pair_exclusion( graph: NeighborGraph, atype: Array, @@ -264,7 +309,8 @@ def apply_pair_exclusion( ------- NeighborGraph A ``dataclasses.replace`` copy (or the original ``graph`` on early - exit) with the exclusion applied. + exit) with the exclusion applied. Existing CSR views are cleared + because their row pointers depend on the valid-edge mask. Notes ----- @@ -285,6 +331,11 @@ def apply_pair_exclusion( out = dataclasses.replace( graph, edge_mask=xp.logical_and(graph.edge_mask, xp.astype(keep, xp.bool)), + destination_order=None, + destination_row_ptr=None, + source_order=None, + source_row_ptr=None, + destination_sorted=False, ) if compact: # Angle fields are a coupled pair (produced together by the angle diff --git a/deepmd/dpmodel/utils/neighbor_graph/pairs.py b/deepmd/dpmodel/utils/neighbor_graph/pairs.py index 75f2682f20..15f45a2ff1 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/pairs.py +++ b/deepmd/dpmodel/utils/neighbor_graph/pairs.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Pairs of edges sharing a center (``dst``) — the edge-pair axis. -Shared primitive: graph-native attention (NeighborGraph PR-D) uses +Graph-native attention uses ``(ordered=True, include_self=True)`` = the full transformer neighbor-pair -square per center; 3-body angles (PR-E) use ``(ordered=False, +square per center; 3-body angles use ``(ordered=False, include_self=False)``. Two forms: @@ -16,8 +16,8 @@ UNBACKED SymInt sizes via :func:`xp_hint_dynamic_size`, so the form traces through ``make_fx``/``torch.export`` and compiles under AOTI (torch >= 2.6 unbacked-symint support) — this is what makes the carry-all attention - graph lower exportable to a ``.pt2``. numpy/jax run it eagerly as before - (jax.jit would still need a static realization — deferred to the jax PR). + graph lower exportable to a ``.pt2``. NumPy and JAX run it eagerly; + ``jax.jit`` requires the shape-static form. - **shape-static** (``static_nnei`` set): assumes the center-major static layout (``E = n_center * static_nnei``, edge ``c * static_nnei + m`` belongs to center ``c`` — the layout ``from_dense_quartet(compact=False)`` emits). diff --git a/deepmd/entrypoints/convert_backend.py b/deepmd/entrypoints/convert_backend.py index e5cd51b386..43cb901449 100644 --- a/deepmd/entrypoints/convert_backend.py +++ b/deepmd/entrypoints/convert_backend.py @@ -36,19 +36,18 @@ def convert_backend( inp_hook = inp_backend.serialize_hook out_hook = out_backend.deserialize_hook data = inp_hook(INPUT) - # Forward atomic_virial to pt_expt deserialize_to_file if applicable; - # warn and skip the flag for backends that don't accept it so that - # scripts passing --atomic-virial indiscriminately don't break. import inspect sig = inspect.signature(out_hook) + hook_kwargs: dict[str, Any] = {} + if "lower_kind" in sig.parameters: + hook_kwargs["lower_kind"] = "auto" if "do_atomic_virial" in sig.parameters: - out_hook(OUTPUT, data, do_atomic_virial=atomic_virial) - else: - if atomic_virial: - log.warning( - "--atomic-virial is only meaningful for pt_expt .pt2/.pte " - "outputs; ignoring it for output backend %s", - out_backend.name, - ) - out_hook(OUTPUT, data) + hook_kwargs["do_atomic_virial"] = atomic_virial + elif atomic_virial: + log.warning( + "--atomic-virial is only meaningful for pt_expt .pt2/.pte " + "outputs; ignoring it for output backend %s", + out_backend.name, + ) + out_hook(OUTPUT, data, **hook_kwargs) diff --git a/deepmd/kernels/autotune.py b/deepmd/kernels/autotune.py new file mode 100644 index 0000000000..a6564294e5 --- /dev/null +++ b/deepmd/kernels/autotune.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Freeze-time Triton launch-configuration autotuning. + +Custom-kernel families register a tuner through :func:`register_autotuner`; a +model-freezing entry point calls :func:`run_autotune` once, before tracing, so +the exported artifact bakes launch configurations tuned for the exact +deployment shapes and GPU. A tuner sweeps only the shape keys its kernels need +that are not already covered by the built-in tables, keeping repeat freezes +cheap. + +Autotuning is a no-op below ``DP_TRITON_INFER`` level 2 (level-1 kernels run a +single shape-independent configuration) and off CUDA -- Triton launch tables +are GPU-specific and AOTInductor artifacts are not portable across GPU models, +so the tuning target is always the local deployment GPU. + +Adding a new custom-kernel family is one registration call at import time; the +freeze paths need no change. +""" + +from __future__ import ( + annotations, +) + +import logging +from collections.abc import ( + Callable, +) + +import torch + +from deepmd.kernels.utils import ( + triton_infer_level, +) + +log = logging.getLogger(__name__) + +# Family name -> tuner(model, level, device). Keying by name makes repeated +# registration (e.g. module re-import) idempotent. +Autotuner = Callable[[torch.nn.Module, int, torch.device], None] +_AUTOTUNERS: dict[str, Autotuner] = {} + + +def register_autotuner(name: str, tuner: Autotuner) -> None: + """Register a kernel family's freeze-time launch-configuration tuner. + + Parameters + ---------- + name : str + Unique family identifier; a repeated name replaces the prior tuner. + tuner : Callable[[torch.nn.Module, int, torch.device], None] + Sweeps the shape keys ``model`` needs for this family that are not + already covered and registers the winners into the family's launch + table. Receives the resolved ``DP_TRITON_INFER`` level and the target + CUDA device. + """ + _AUTOTUNERS[name] = tuner + + +def run_autotune(model: torch.nn.Module, device: torch.device) -> None: + """Tune every registered kernel family for ``model`` on ``device``. + + Called by a freeze path after the model is built and before it is traced, + so ``resolve_*_config`` lookups made while tracing return freshly tuned + launches. A no-op below ``DP_TRITON_INFER`` level 2 or off CUDA. + + Parameters + ---------- + model : torch.nn.Module + The model about to be frozen; tuners walk it for their shape keys. + device : torch.device + The target CUDA device the artifact will run on. + """ + level = triton_infer_level() + if level < 2 or device.type != "cuda" or not torch.cuda.is_available(): + return + for name, tuner in _AUTOTUNERS.items(): + # A tuning failure must degrade to the built-in / default launch tables + # rather than abort the freeze; the frozen artifact stays correct, only + # its launch configuration is untuned for uncovered shapes. + try: + tuner(model, level, device) + except Exception: + log.warning( + "Triton autotuner %r failed; keeping default launch configurations.", + name, + exc_info=True, + ) diff --git a/deepmd/kernels/cuda/__init__.py b/deepmd/kernels/cuda/__init__.py new file mode 100644 index 0000000000..2ce6f20cc1 --- /dev/null +++ b/deepmd/kernels/cuda/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Hand-written CUDA / cuBLAS operators for graph-lower inference. + +The CUDA sources live under ``source/op/pt`` and compile into +``libdeepmd_op_pt.so`` (loaded via ``deepmd.pt.cxx_op``); the modules here +expose the resulting ``torch.ops.deepmd.*`` operators to the pt_expt graph +lower together with the backward, meta (fake) and CPU trace-time +implementations that ``torch.export`` / ``make_fx`` require. Dispatch is +gated by ``DP_CUDA_INFER`` (:func:`deepmd.kernels.utils.cuda_infer_level`). + +Modules +------- +:mod:`.dpa1.graph_descriptor` + DPA1 (``se_atten``) descriptor mega kernels: environment matrix, + embedding MLP, moment reduction and ``G^T G`` contraction in one forward + / 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). +:mod:`.edge_force_virial` + Descriptor-agnostic force / atom-virial / per-frame-virial assembly from + the per-edge energy gradient. +""" diff --git a/deepmd/kernels/cuda/dpa1/__init__.py b/deepmd/kernels/cuda/dpa1/__init__.py new file mode 100644 index 0000000000..7d394ac032 --- /dev/null +++ b/deepmd/kernels/cuda/dpa1/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Fused CUDA descriptors for the DPA1 (``se_atten``) graph lower. + +:mod:`.graph_descriptor` binds the ``deepmd::dpa1_graph_descriptor`` mega +kernels (environment matrix, three-layer embedding MLP, moment and ``G^T G``). +:mod:`.graph_compress` binds ``deepmd::dpa1_graph_compress`` for the +geo-compressed strip path (quintic table lookup plus the precomputed strip +type-pair gate). Dispatch is gated by :meth:`DescrptDPA1._fused_eligible` at +``DP_CUDA_INFER >= 1``. +""" diff --git a/deepmd/kernels/cuda/dpa1/canonical.py b/deepmd/kernels/cuda/dpa1/canonical.py new file mode 100644 index 0000000000..8af278caa7 --- /dev/null +++ b/deepmd/kernels/cuda/dpa1/canonical.py @@ -0,0 +1,454 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""CUDA bindings for compact canonical compressed DPA1 deployment.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch + +if TYPE_CHECKING: + from deepmd.pt_expt.utils.canonical_graph import ( + DPA1CanonicalGraph, + ) + +__all__ = [ + "canonical_model_eligible", + "dpa1_canonical_compress_energy_force", + "ensure_registered", + "op_available", +] + +_cpu_library: torch.library.Library | None = None + + +def canonical_model_eligible(model: Any) -> bool: + """Return whether a model can use the no-mask compact deployment 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 + if not bool(getattr(descriptor, "geo_compress", False)): + return False + eligible = getattr(descriptor, "_fused_eligible", None) + if not callable(eligible) or not bool(eligible("cuda")): + 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.dpa1.graph_compress import ( + mega_eligible, + ) + from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + ) + + return mega_eligible(descriptor) and fitting_eligible(fitting) + + +def op_available() -> bool: + """Return whether both compact canonical descriptor operators are loaded.""" + forward = getattr(torch.ops.deepmd, "dpa1_canonical_compress", None) + backward = getattr(torch.ops.deepmd, "dpa1_canonical_compress_backward", None) + return isinstance(forward, torch._ops.OpOverloadPacket) and isinstance( + backward, torch._ops.OpOverloadPacket + ) + + +def _forward_fake( + edge_vec: torch.Tensor, + source: torch.Tensor, + destination_row_ptr: torch.Tensor, + atype: torch.Tensor, + type_embedding: torch.Tensor, + average: torch.Tensor, + inverse_stddev: torch.Tensor, + table: torch.Tensor, + gate_table: torch.Tensor, + type_one_side: int, + concatenate_type_embedding: int, + write_rotation: int, + smooth: int, + axis: int, + lower: float, + upper: float, + table_max: float, + stride0: float, + stride1: float, + rcut: float, + rcut_smooth: float, + protection: float, + neighbors: float, +) -> tuple[torch.Tensor, ...]: + del ( + source, + destination_row_ptr, + average, + inverse_stddev, + gate_table, + type_one_side, + smooth, + lower, + upper, + table_max, + stride0, + stride1, + rcut, + rcut_smooth, + protection, + neighbors, + ) + node_count = atype.shape[0] + width = table.shape[1] // 6 + output_width = width * axis + ( + type_embedding.shape[1] if concatenate_type_embedding else 0 + ) + return ( + edge_vec.new_empty(node_count, output_width), + edge_vec.new_empty( + node_count if write_rotation else 0, + width, + 3, + ), + edge_vec.new_empty(node_count, 4, width), + ) + + +def _backward_fake( + descriptor_gradient: torch.Tensor, + rotation_gradient: torch.Tensor | None, + moment: torch.Tensor, + edge_vec: torch.Tensor, + source: torch.Tensor, + destination_row_ptr: torch.Tensor, + atype: torch.Tensor, + average: torch.Tensor, + inverse_stddev: torch.Tensor, + table: torch.Tensor, + gate_table: torch.Tensor, + type_one_side: int, + smooth: int, + axis: int, + lower: float, + upper: float, + table_max: float, + stride0: float, + stride1: float, + rcut: float, + rcut_smooth: float, + protection: float, + neighbors: float, +) -> torch.Tensor: + del ( + descriptor_gradient, + rotation_gradient, + moment, + source, + destination_row_ptr, + atype, + average, + inverse_stddev, + table, + gate_table, + type_one_side, + smooth, + axis, + lower, + upper, + table_max, + stride0, + stride1, + rcut, + rcut_smooth, + protection, + neighbors, + ) + return torch.empty_like(edge_vec) + + +def _generic_topology( + source: torch.Tensor, + destination_row_ptr: torch.Tensor, +) -> tuple[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( + 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=source.dtype, + device=source.device, + ) + return edge_index, edge_mask, destination_order + + +def _cpu_forward(*args: Any) -> tuple[torch.Tensor, ...]: + from deepmd.kernels.cuda.dpa1.graph_compress import _cpu_forward as generic_forward + + edge_vec, source, destination_row_ptr, *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, + *tail[:11], + True, + *tail[11:], + ) + + +def _cpu_backward(*args: Any) -> torch.Tensor: + from deepmd.kernels.cuda.dpa1.graph_compress import ( + _cpu_backward as generic_backward, + ) + + descriptor_gradient, rotation_gradient, moment, edge_vec = args[:4] + source, destination_row_ptr = args[4:6] + tail = args[6:] + edge_index, edge_mask, destination_order = _generic_topology( + source, + destination_row_ptr, + ) + return generic_backward( + descriptor_gradient, + rotation_gradient, + moment, + edge_vec, + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + *tail[:8], + True, + *tail[8:], + ) + + +def ensure_registered() -> None: + """Register fake and CPU implementations for compact canonical operators.""" + global _cpu_library + if _cpu_library is not None or not op_available(): + return + torch.library.register_fake("deepmd::dpa1_canonical_compress")(_forward_fake) + torch.library.register_fake("deepmd::dpa1_canonical_compress_backward")( + _backward_fake + ) + _cpu_library = torch.library.Library("deepmd", "IMPL") + _cpu_library.impl("dpa1_canonical_compress", _cpu_forward, "CPU") + _cpu_library.impl( + "dpa1_canonical_compress_backward", + _cpu_backward, + "CPU", + ) + + +def dpa1_canonical_compress_energy_force( + desc: Any, + fit: Any, + graph: DPA1CanonicalGraph, + atype: torch.Tensor, + type_embedding: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate compressed DPA1 from a compact canonical edge stream. + + Parameters + ---------- + desc + Geometrically compressed DPA1 descriptor. + fit + Compatible mixed-type energy fitting network. + graph + Compact destination-major graph. + atype + Flat node atom types with shape ``(N,)``. + type_embedding + Type embedding table. + ownership + Energy-contributing node mask with shape ``(N,)``. + atom_bias + Combined fitting and atomic-model energy bias. + do_atomic_virial + Whether to materialize per-node virial. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + Frame energy, atom energy, force, frame virial, and atom virial. + """ + from deepmd.kernels.cuda.dpa1.graph_compress import ( + 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.graph_fitting import ( + ensure_registered as ensure_fitting_registered, + ) + + ensure_registered() + ensure_fitting_registered() + ensure_force_registered() + if not mega_eligible(desc) or not canonical_op_available(): + raise ValueError("model is not eligible for compact canonical DPA1 inference") + + se = desc.se_atten + compress_data = desc.compress_data[0].contiguous() + gate_table = desc.type_embd_data.contiguous() + inverse_stddev = torch.reciprocal(se.stddev[:, 0, :]).contiguous() + from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ) + + with disable_proxy_modes_tracing(): + lower, upper, table_max, stride0, stride1 = ( + float(value) for value in desc.compress_info[0].tolist()[:5] + ) + + descriptor, _rotation, moment = torch.ops.deepmd.dpa1_canonical_compress( + graph.edge_vec, + graph.source, + graph.destination_row_ptr, + atype, + type_embedding, + se.mean[:, 0, :].contiguous(), + inverse_stddev, + compress_data, + gate_table, + int(se.type_one_side), + int(desc.concat_output_tebd), + 0, + int(se.smooth), + int(se.axis_neuron), + lower, + upper, + table_max, + stride0, + stride1, + float(se.rcut), + float(se.rcut_smth), + float(se.env_protection), + float(se.nnei), + ) + + *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, + ) + + 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 + ), + atom_bias.to(torch.float64).contiguous(), + ACT_CODES[str(hidden[0].activation_function).lower()], + ) + 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, + ) + + 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) + del descriptor + descriptor_gradient = torch.ops.deepmd.graph_fitting_backward( + energy_seed, + fitting_saved, + weights, + residuals, + head.w.reshape(-1).contiguous(), + ) + del fitting_saved + edge_gradient = torch.ops.deepmd.dpa1_canonical_compress_backward( + descriptor_gradient, + None, + moment, + graph.edge_vec, + graph.source, + graph.destination_row_ptr, + atype, + se.mean[:, 0, :].contiguous(), + inverse_stddev, + compress_data, + gate_table, + int(se.type_one_side), + int(se.smooth), + int(se.axis_neuron), + lower, + upper, + table_max, + stride0, + stride1, + float(se.rcut), + float(se.rcut_smth), + float(se.env_protection), + float(se.nnei), + ) + force, atom_virial, virial = canonical_edge_force_virial( + edge_gradient, + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + graph.n_node, + atype.shape[0], + do_atomic_virial, + ) + return energy, atom_energy, force, virial, atom_virial diff --git a/deepmd/kernels/cuda/dpa1/graph_compress.py b/deepmd/kernels/cuda/dpa1/graph_compress.py new file mode 100644 index 0000000000..1bad0501d1 --- /dev/null +++ b/deepmd/kernels/cuda/dpa1/graph_compress.py @@ -0,0 +1,992 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the fused DPA1 compressed graph-lower descriptor operator. + +The CUDA operator ``deepmd::dpa1_graph_compress`` (see +``source/op/pt/dpa1_graph_compress.cu``) evaluates the geo-compressed, +attention-free strip descriptor from the flat edge stream -- environment matrix, +quintic spline table lookup of the geometric embedding, the precomputed strip +type-pair gate, moment reduction and the ``G^T G`` contraction -- in two mega +kernels (forward and backward). This module mirrors +:mod:`.graph_descriptor` for the compressed path: + +* **Fake (meta) implementations** for ``make_fx`` / ``torch.export``. +* **A registered backward** routing to ``deepmd::dpa1_graph_compress_backward``. +* **CPU reference implementations** for the freeze pipeline's trace-time + sample evaluation. + +The ``gate_table`` argument carries the precomputed strip type-pair embedding +(:attr:`DescrptDPA1.type_embd_data`); ``compress_data`` / ``compress_info`` hold +the quintic spline coefficients and segment metadata of the tabulated geometric +net. +""" + +from typing import ( + Any, +) + +import torch + +__all__ = [ + "dpa1_graph_compress", + "dpa1_graph_compress_energy_force", + "ef_op_available", + "ensure_registered", + "mega_eligible", + "op_available", +] + +# Instantiated table widths in dpa1_graph_compress.cu; each divides the 256 +# threads per block, so the moment walk splits channels evenly. An arbitrary +# ``ng`` is served by the smallest width that is at least ``ng``, zero-padding +# the spline table and the type-pair gate so the extra channels contribute +# nothing, and slicing them off the descriptor. +_KERNEL_WIDTHS = (8, 16, 32, 64, 128, 256) + + +def _bucket_width(ng: int) -> int: + """Smallest instantiated kernel width that holds ``ng`` channels.""" + for w in _KERNEL_WIDTHS: + if w >= ng: + return w + raise ValueError( + f"dpa1_graph_compress: ng={ng} exceeds the maximum table width " + f"{_KERNEL_WIDTHS[-1]}" + ) + + +def op_available() -> bool: + """Whether the C++ ``deepmd::dpa1_graph_compress`` op is loaded.""" + op = getattr(torch.ops.deepmd, "dpa1_graph_compress", None) + return isinstance(op, torch._ops.OpOverloadPacket) + + +def ef_op_available() -> bool: + """Whether all operators required by the fused inference assembly 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(desc: Any) -> bool: + """Whether ``desc`` can use the fused energy-force mega operator. + + The mega operator feeds the tabulated descriptor straight into the energy + fitting, so the descriptor width the kernel emits must equal the fitting + input width. This holds only when ``ng`` is an instantiated kernel width + (no bucket zero-padding). A non-bucket ``ng`` stays on the level-1 lower, + where the padding channels are sliced off before the fitting. + """ + ng = int(desc.se_atten.neuron[-1]) + return _bucket_width(ng) == ng + + +# ====================================================================== +# Fake (meta) implementations +# ====================================================================== +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, + type_embedding: torch.Tensor, + davg: torch.Tensor, + inverse_stddev: torch.Tensor, + table: torch.Tensor, + gate_table: torch.Tensor, + type_one_side: int, + concat_tebd: int, + write_rotation: int, + smooth: int, + axis: int, + canonical: bool, + lower: float, + upper: float, + table_max: float, + stride0: float, + stride1: float, + rcut: float, + rcut_smth: float, + protection: float, + nnei: float, +) -> tuple[torch.Tensor, ...]: + n_node = atype.shape[0] + ng = table.shape[1] // 6 + _, tebd_dim = type_embedding.shape + out_dim = ng * axis + (tebd_dim if concat_tebd else 0) + dev = edge_vec.device + return ( + torch.empty(n_node, out_dim, dtype=torch.float32, device=dev), + torch.empty( + n_node if write_rotation else 0, + ng, + 3, + dtype=torch.float32, + device=dev, + ), + torch.empty(n_node, 4, ng, dtype=torch.float32, device=dev), + ) + + +def _backward_fake( + d_grrg: torch.Tensor, + d_rot_mat: torch.Tensor | None, + gr: torch.Tensor, + edge_vec: torch.Tensor, + *args: Any, +) -> torch.Tensor: + return torch.empty_like(edge_vec) + + +# ====================================================================== +# Autograd bridge +# ====================================================================== +def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: + ( + edge_vec, + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + atype, + type_embedding, + davg, + inverse_stddev, + table, + gate_table, + type_one_side, + concat_tebd, + _write_rotation, + smooth, + axis, + canonical, + lower, + upper, + table_max, + stride0, + stride1, + rcut, + rcut_smth, + protection, + nnei, + ) = inputs + (gr,) = output[2:] + ctx.save_for_backward( + gr, + edge_vec, + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + atype, + davg, + inverse_stddev, + table, + gate_table, + ) + ctx.scalars = ( + type_one_side, + concat_tebd, + smooth, + axis, + canonical, + lower, + upper, + table_max, + stride0, + stride1, + rcut, + rcut_smth, + protection, + nnei, + ) + ctx.set_materialize_grads(False) + + +def _backward( + ctx: Any, + d_grrg: torch.Tensor, + d_rot_mat: torch.Tensor | None, + *d_aux: Any, +) -> tuple: + ( + gr, + edge_vec, + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + atype, + davg, + inverse_stddev, + table, + gate_table, + ) = ctx.saved_tensors + ( + type_one_side, + concat_tebd, + smooth, + axis, + canonical, + lower, + upper, + table_max, + stride0, + stride1, + rcut, + rcut_smth, + protection, + nnei, + ) = ctx.scalars + d_edge_vec = torch.ops.deepmd.dpa1_graph_compress_backward( + d_grrg, + d_rot_mat, + gr, + edge_vec, + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + atype, + davg, + inverse_stddev, + table, + gate_table, + type_one_side, + smooth, + axis, + canonical, + lower, + upper, + table_max, + stride0, + stride1, + rcut, + rcut_smth, + protection, + nnei, + ) + return (d_edge_vec,) + (None,) * 25 + + +# ====================================================================== +# CPU reference implementations +# ====================================================================== +def _cpu_tabulate_fusion( + table: torch.Tensor, + lower: float, + upper: float, + table_max: float, + stride0: float, + stride1: float, + em_x: torch.Tensor, + last_layer_size: int, + reference: torch.Tensor, +) -> torch.Tensor: + """Quintic spline table lookup with C1 extrapolation (fp32). + + Parameters + ---------- + table + Coefficient table with shape (nspline, last_layer_size * 6). + lower, upper, table_max, stride0, stride1 + Spline segment bounds and strides (``compress_info`` scalars). + em_x + Lookup coordinates with shape (nloc, nnei). + last_layer_size + Output width ``ng``. + reference + Tensor whose dtype/device define the compute context. + + Returns + ------- + torch.Tensor + Tabulated values with shape (nloc, nnei, last_layer_size). + """ + nloc, nnei = em_x.shape[:2] + xx = em_x.reshape(nloc, nnei) + tbl = table.to(device=reference.device, dtype=reference.dtype) + + zeros = torch.zeros(xx.shape, dtype=torch.int64, device=xx.device) + nspline = tbl.shape[0] + last_idx = torch.full(xx.shape, nspline - 1, dtype=torch.int64, device=xx.device) + first_stride = int((upper - lower) / stride0) + first_stride_value = float(first_stride) + + first_idx = torch.floor((xx - lower) / stride0).to(torch.int64) + second_idx = first_stride + torch.floor((xx - upper) / stride1).to(torch.int64) + table_idx = torch.where( + xx < lower, + zeros, + torch.where( + xx < upper, + first_idx, + torch.where(xx < table_max, second_idx, last_idx), + ), + ) + table_idx = torch.clamp(table_idx, min=0, max=nspline - 1) + + table_idx_value = table_idx.to(reference.dtype) + dx_first = xx - (table_idx_value * stride0 + lower) + dx_second = xx - ((table_idx_value - first_stride_value) * stride1 + upper) + dx_high = table_max - ( + (last_idx.to(reference.dtype) - first_stride_value) * stride1 + upper + ) + dx = torch.where( + xx < lower, + torch.zeros_like(xx), + torch.where( + xx < upper, + dx_first, + torch.where(xx < table_max, dx_second, dx_high), + ), + ) + extrapolate_delta = torch.where( + xx < lower, + xx - lower, + torch.where(xx >= table_max, xx - table_max, torch.zeros_like(xx)), + ) + + coeff = tbl[table_idx.reshape(-1)] + coeff = coeff.reshape(nloc, nnei, last_layer_size, 6) + dx = dx.reshape(nloc, nnei, 1) + values = ( + coeff[..., 0] + + ( + coeff[..., 1] + + ( + coeff[..., 2] + + (coeff[..., 3] + (coeff[..., 4] + coeff[..., 5] * dx) * dx) * dx + ) + * dx + ) + * dx + ) + values_grad = ( + coeff[..., 1] + + ( + 2 * coeff[..., 2] + + (3 * coeff[..., 3] + (4 * coeff[..., 4] + 5 * coeff[..., 5] * dx) * dx) + * dx + ) + * dx + ) + extrapolate_delta = extrapolate_delta.reshape(nloc, nnei, 1) + return values + values_grad * extrapolate_delta + + +def _cpu_tabulate_se_atten( + table: torch.Tensor, + info: tuple[float, float, float, float, float], + em_x: torch.Tensor, + em: torch.Tensor, + two_embed: torch.Tensor, + last_layer_size: int, +) -> torch.Tensor: + """Per-edge tabulate_fusion_se_atten body returning the moment outer factor. + + Returns ``(E, 4, ng)`` with ``em`` shaped ``(E, 1, 4)``. + """ + values = _cpu_tabulate_fusion(table, *info, em_x, last_layer_size, em) + values = values * two_embed + values + return (em.unsqueeze(-1) * values.unsqueeze(2)).sum(dim=1) + + +def _cpu_pair_idx( + edge_index: torch.Tensor, + atype: torch.Tensor, + type_one_side: int, + ntypes: int, +) -> torch.Tensor: + src, dst = edge_index[0], edge_index[1] + nei_type, center_type = atype[src], atype[dst] + if type_one_side: + return nei_type + return center_type * ntypes + nei_type + + +def _cpu_env_and_gg( + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + atype: torch.Tensor, + gate_table: torch.Tensor, + davg: torch.Tensor, + dstd: torch.Tensor, + table: torch.Tensor, + info: tuple[float, float, float, float, float], + type_one_side: int, + smooth: int, + ntypes: int, + ng: int, + rcut: float, + rcut_smth: float, + protection: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Environment matrix, tabulated geometric net and strip gate (fp32). + + Returns ``(rr, outer)`` where ``outer`` is ``(E, 4, ng)`` before the + neighbor-axis reduction. + """ + ev = edge_vec.to(torch.float32) + src, dst = edge_index[0], edge_index[1] + center_type = atype[dst] + length = ev.norm(dim=-1, keepdim=True) + length = length + (~edge_mask[:, None]).to(length.dtype) + q = length + protection + u = ((length - rcut_smth) / (rcut - rcut_smth)).clamp(0.0, 1.0) + sw = u**3 * (-6 * u**2 + 15 * u - 10) + 1.0 + em = torch.cat([sw / q, ev * (sw / q**2)], dim=-1) + rr = (em - davg[center_type]) / dstd[center_type] + ss = rr[:, 0:1] + pair_idx = _cpu_pair_idx(edge_index, atype, type_one_side, ntypes) + gate = gate_table[pair_idx] + if smooth: + gate = gate * sw + n_edge = edge_vec.shape[0] + em_x = ss.reshape(n_edge, 1) + em_rr = rr.reshape(n_edge, 1, 4) + two_embed = gate.reshape(n_edge, 1, ng) + outer = _cpu_tabulate_se_atten(table, info, em_x, em_rr, two_embed, ng) + return rr, outer + + +def _cpu_outputs( + rr: torch.Tensor, + outer: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + atype: torch.Tensor, + type_embedding: torch.Tensor | None, + ng: int, + axis: int, + concat_tebd: int, + nnei: float, + n_node: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + outer = outer * edge_mask[:, None, None].to(outer.dtype) + gr = torch.zeros(n_node, 4, ng, dtype=outer.dtype, device=outer.device) + gr.index_add_(0, edge_index[1], outer) + gr = gr / nnei + gr_t = gr.permute(0, 2, 1) + grrg = torch.matmul(gr_t, gr[:, :, :axis]).reshape(n_node, ng * axis) + rot_mat = gr_t[:, :, 1:4].contiguous() + if concat_tebd: + grrg = torch.cat([grrg, type_embedding[atype]], dim=-1) + return grrg.contiguous(), rot_mat, gr + + +def _cpu_forward( + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + atype: torch.Tensor, + type_embedding: torch.Tensor, + davg: torch.Tensor, + inverse_stddev: torch.Tensor, + table: torch.Tensor, + gate_table: torch.Tensor, + type_one_side: int, + concat_tebd: int, + write_rotation: int, + smooth: int, + axis: int, + canonical: bool, + lower: float, + upper: float, + table_max: float, + stride0: float, + stride1: float, + rcut: float, + rcut_smth: float, + protection: float, + nnei: float, +) -> tuple[torch.Tensor, ...]: + n_node = atype.shape[0] + ng = table.shape[1] // 6 + ntypes = type_embedding.shape[0] + if n_node == 0: + out_dim = ng * axis + (type_embedding.shape[1] if concat_tebd else 0) + return ( + table.new_empty(0, out_dim), + table.new_empty(0, ng, 3), + table.new_empty(0, 4, ng), + ) + rr, outer = _cpu_env_and_gg( + edge_vec, + edge_index, + edge_mask, + atype, + gate_table, + davg, + torch.reciprocal(inverse_stddev), + table, + (lower, upper, table_max, stride0, stride1), + type_one_side, + smooth, + ntypes, + ng, + rcut, + rcut_smth, + protection, + ) + grrg, rot_mat, gr = _cpu_outputs( + rr, + outer, + edge_index, + edge_mask, + atype, + type_embedding, + ng, + axis, + concat_tebd, + nnei, + n_node, + ) + if not write_rotation: + rot_mat = rot_mat.new_empty(0, ng, 3) + return grrg, rot_mat, gr + + +def _cpu_backward( + d_grrg: torch.Tensor, + d_rot_mat: torch.Tensor | None, + gr: 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, + davg: torch.Tensor, + inverse_stddev: torch.Tensor, + table: torch.Tensor, + gate_table: torch.Tensor, + type_one_side: int, + smooth: int, + axis: int, + canonical: bool, + lower: float, + upper: float, + table_max: float, + stride0: float, + stride1: float, + rcut: float, + rcut_smth: float, + protection: float, + nnei: float, +) -> torch.Tensor: + n_node = atype.shape[0] + ng = table.shape[1] // 6 + if n_node == 0: + return torch.zeros_like(edge_vec) + ntypes = gate_table.shape[0] if type_one_side else round(gate_table.shape[0] ** 0.5) + ev = edge_vec.detach().clone().requires_grad_(True) + with torch.enable_grad(): + rr, outer = _cpu_env_and_gg( + ev, + edge_index, + edge_mask, + atype, + gate_table, + davg, + torch.reciprocal(inverse_stddev), + table, + (lower, upper, table_max, stride0, stride1), + type_one_side, + smooth, + ntypes, + ng, + rcut, + rcut_smth, + protection, + ) + grrg, rot_mat, _gr = _cpu_outputs( + rr, + outer, + edge_index, + edge_mask, + atype, + None, + ng, + axis, + 0, + nnei, + n_node, + ) + loss = (grrg * d_grrg[:, : grrg.shape[1]].to(grrg.dtype)).sum() + if d_rot_mat is not None: + loss = loss + (rot_mat * d_rot_mat.to(rot_mat.dtype)).sum() + (d_ev,) = torch.autograd.grad(loss, ev) + return d_ev.to(edge_vec.dtype) + + +# ====================================================================== +# Registration and the public wrapper +# ====================================================================== +_cpu_library: torch.library.Library | None = None + + +def ensure_registered() -> None: + """Register the fake / backward / CPU implementations for the ops. + + Idempotent; a no-op when the C++ operator library is not loaded. + """ + global _cpu_library + if _cpu_library is not None or not op_available(): + return + torch.library.register_fake("deepmd::dpa1_graph_compress")(_forward_fake) + torch.library.register_fake("deepmd::dpa1_graph_compress_backward")(_backward_fake) + torch.library.register_autograd( + "deepmd::dpa1_graph_compress", _backward, setup_context=_setup_context + ) + _cpu_library = torch.library.Library("deepmd", "IMPL") + _cpu_library.impl("dpa1_graph_compress", _cpu_forward, "CPU") + _cpu_library.impl("dpa1_graph_compress_backward", _cpu_backward, "CPU") + + +def dpa1_graph_compress( + desc: Any, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Evaluate the fused geo-compressed DPA1 strip descriptor from the edge stream. + + Numerically equivalent to the dense :meth:`DescrptDPA1._call_compressed` + body on the attention-free strip configuration (up to fp32 summation order). + ``edge_vec`` may arrive fp64; the operator computes in fp32 and the backward + returns ``d_edge_vec`` in the leaf dtype. + + Parameters + ---------- + desc : DescrptDPA1 + The pt_expt descriptor module with ``compress`` and ``geo_compress`` + enabled; must satisfy ``desc._fused_eligible("cuda")``. + graph : NeighborGraph + Lowered neighbor graph with ``edge_vec`` (E, 3), ``edge_index`` (2, E) + and ``edge_mask`` (E,), plus a destination CSR view. The CSR may cover + cached skin edges whose current cutoff mask is false. + atype : torch.Tensor + Flat node atom types with shape (N,), int64. + type_embedding : torch.Tensor + Type embedding table with shape (ntypes + 1, tebd_dim), fp32. + + Returns + ------- + grrg : torch.Tensor + Descriptor with shape (N, ng * axis [+ tebd_dim]), fp32. + rot_mat : torch.Tensor + Equivariant rotation matrix with shape (N, ng, 3), fp32. + """ + ensure_registered() + se = desc.se_atten + ng = int(se.neuron[-1]) + axis = int(se.axis_neuron) + # The kernel is instantiated at a fixed set of table widths; a descriptor + # whose ``ng`` is not one of them runs at the next larger width with the + # spline table and type-pair gate zero-padded, so the extra channels are + # identically zero and are sliced off below. + width = _bucket_width(ng) + pad = width - ng + compress_data = desc.compress_data[0].contiguous() + gate_table = desc.type_embd_data.contiguous() + if pad: + compress_data = torch.nn.functional.pad(compress_data, (0, pad * 6)) + gate_table = torch.nn.functional.pad(gate_table, (0, pad)) + # The spline segment bounds / strides are baked as scalar constants (not a + # host tensor) so the exported graph carries no CPU-resident operand. The + # compress-info buffer is a graph constant; its values are read with proxy + # tracing disabled so make_fx bakes plain floats rather than proxying the + # scalar extraction. + from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ) + + with disable_proxy_modes_tracing(): + lower, upper, table_max, stride0, stride1 = ( + float(x) for x in desc.compress_info[0].tolist()[:5] + ) + if graph.destination_order is None or graph.destination_row_ptr is None: + raise ValueError("dpa1_graph_compress requires destination CSR topology") + grrg, rot_mat, _moment = torch.ops.deepmd.dpa1_graph_compress( + graph.edge_vec.contiguous(), + graph.edge_index.contiguous(), + graph.edge_mask.contiguous(), + graph.destination_order.contiguous(), + graph.destination_row_ptr.contiguous(), + atype.contiguous(), + type_embedding.contiguous(), + se.mean[:, 0, :].contiguous(), + torch.reciprocal(se.stddev[:, 0, :]).contiguous(), + compress_data, + gate_table, + int(se.type_one_side), + int(desc.concat_output_tebd), + 1, + int(se.smooth), + int(se.axis_neuron), + bool(graph.destination_sorted), + lower, + upper, + table_max, + stride0, + stride1, + float(se.rcut), + float(se.rcut_smth), + float(se.env_protection), + float(se.nnei), + ) + if pad: + # Drop the padding channels: the descriptor is stored channel-major + # (i * axis + j, i in [0, width)), so the real block is the first + # ng * axis columns; any concatenated type-embedding tail follows the + # full width * axis span. The rotation matrix keeps its first ng rows. + desc_block = grrg[:, : ng * axis] + if desc.concat_output_tebd: + grrg = torch.cat([desc_block, grrg[:, width * axis :]], dim=-1) + else: + grrg = desc_block + rot_mat = rot_mat[:, :ng, :] + return grrg, rot_mat + + +def dpa1_graph_compress_energy_force( + desc: Any, + fit: Any, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + node_capacity: int, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate compressed DPA1 energy, force, and virial without autograd. + + The descriptor forward/backward and force assembly remain explicit custom + operators in the exported graph. The fitting forward/backward is evaluated + between them, so force is returned as a value and no autograd tape is + retained. Destination/source CSR tensors are part of the graph contract and + are reused by both descriptor and force kernels. A graph marked + ``destination_sorted`` selects direct destination-major addressing; edge + masks remain authoritative in both addressing modes. This inference-only + path suppresses the unused rotation output, and fitting backward does not + retain the descriptor after its last value use. + + Parameters + ---------- + desc : DescrptDPA1 + The compressed pt_expt descriptor; must satisfy ``mega_eligible(desc)`` + and ``desc._fused_eligible("cuda")``. + fit : EnergyFittingNet + The pt_expt fitting module (see + :func:`~deepmd.kernels.cuda.graph_fitting.fitting_eligible`). + graph : NeighborGraph + The lowered neighbor graph (``edge_vec``, ``edge_index``, ``edge_mask``, + ``n_node``) with destination/source CSR. ``destination_sorted`` must be + true only when ``destination_order`` is the identity permutation. + atype : torch.Tensor + Flat node atom types with shape (N,), int64. + type_embedding : torch.Tensor + Type embedding table with shape (ntypes + 1, tebd_dim), fp32. + ownership : torch.Tensor + Energy-contributing node mask with shape (N,), bool. + atom_bias : torch.Tensor + Combined fitting and atomic-model bias with shape (ntypes,). + node_capacity : int + Padded node-axis size ``N`` (the force / atom-virial scatter bound). + do_atomic_virial : bool + Whether to also assemble the per-atom virial. + + Returns + ------- + energy : torch.Tensor + Per-frame energy with shape (nf, 1), fp64. + atom_energy : torch.Tensor + Per-atom energy with shape (N, 1), fp64. + force : torch.Tensor + Per-atom force with shape (N, 3), fp32. + virial : torch.Tensor + Per-frame virial with shape (nf, 3, 3), fp32. + atom_virial : torch.Tensor + Per-atom virial with shape (N, 3, 3) when requested, else empty (0, 3, 3). + """ + 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( + "compressed DPA1 inference requires destination/source CSR topology" + ) + se = desc.se_atten + compress_data = desc.compress_data[0].contiguous() + gate_table = desc.type_embd_data.contiguous() + edge_vec = graph.edge_vec.to(compress_data.dtype).contiguous() + # Scalar spline metadata baked as constants (see dpa1_graph_compress), read + # with proxy tracing disabled so make_fx bakes plain floats. + from torch.fx.experimental.proxy_tensor import ( + disable_proxy_modes_tracing, + ) + + with disable_proxy_modes_tracing(): + 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() + descriptor, _rotation, moment = torch.ops.deepmd.dpa1_graph_compress( + edge_vec, + graph.edge_index.contiguous(), + graph.edge_mask.contiguous(), + graph.destination_order.contiguous(), + graph.destination_row_ptr.contiguous(), + atype.contiguous(), + type_embedding.contiguous(), + se.mean[:, 0, :].contiguous(), + inverse_stddev, + compress_data, + gate_table, + int(se.type_one_side), + int(desc.concat_output_tebd), + 0, + int(se.smooth), + int(se.axis_neuron), + bool(graph.destination_sorted), + lower, + upper, + table_max, + stride0, + stride1, + float(se.rcut), + float(se.rcut_smth), + float(se.env_protection), + float(se.nnei), + ) + 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, + ) + + activation = ACT_CODES[str(hidden[0].activation_function).lower()] + atom_energy_raw, fitting_saved = torch.ops.deepmd.graph_fitting( + descriptor, + atype, + weights, + biases, + timesteps, + residuals, + head_weight, + head_bias, + atom_bias, + 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, + ) + + 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) + del descriptor + descriptor_gradient = torch.ops.deepmd.graph_fitting_backward( + energy_seed, + fitting_saved, + weights, + residuals, + head_weight, + ) + del fitting_saved + edge_gradient = torch.ops.deepmd.dpa1_graph_compress_backward( + descriptor_gradient, + None, + moment, + edge_vec, + graph.edge_index, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + atype, + se.mean[:, 0, :].contiguous(), + inverse_stddev, + compress_data, + gate_table, + int(se.type_one_side), + int(se.smooth), + int(se.axis_neuron), + bool(graph.destination_sorted), + lower, + upper, + table_max, + stride0, + stride1, + float(se.rcut), + float(se.rcut_smth), + float(se.env_protection), + float(se.nnei), + ) + force, atom_virial, virial = edge_force_virial( + edge_gradient, + edge_vec, + graph.edge_index, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + graph.n_node, + node_capacity, + do_atomic_virial, + ) + return energy, atom_energy, force, virial, atom_virial diff --git a/deepmd/kernels/cuda/dpa1/graph_descriptor.py b/deepmd/kernels/cuda/dpa1/graph_descriptor.py new file mode 100644 index 0000000000..8d1ff2f486 --- /dev/null +++ b/deepmd/kernels/cuda/dpa1/graph_descriptor.py @@ -0,0 +1,760 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the fused DPA1 graph-lower descriptor operator. + +The CUDA operator ``deepmd::dpa1_graph_descriptor`` (see +``source/op/pt/dpa1_graph_descriptor.cu``) evaluates the whole DPA1 +(``se_atten``) attention-free descriptor body from the flat edge stream -- +environment matrix, three-layer embedding MLP, the strip type-pair gate when +applicable, moment reduction and the ``G^T G`` contraction -- in two mega +kernels (forward and backward). This module supplies everything the pt_expt +graph lower needs beyond the raw CUDA implementation: + +* **Fake (meta) implementations** so ``make_fx`` / ``torch.export`` can infer + output shapes and bake the operator into the graph-form ``.pt2`` without + running CUDA. +* **A registered backward** routing to + ``deepmd::dpa1_graph_descriptor_backward``. The graph lower assembles force + and virial analytically from ``grad(E, edge_vec)``, so the descriptor must + expose an ``edge_vec`` gradient; weights and integer topology are constants + of the inference graph and receive ``None``. +* **CPU reference implementations** registered on the same ops. The pt_expt + freeze pipeline traces and sample-evaluates the model on CPU; these keep the + ops functional there (trace-time only -- deployment always dispatches the + CUDA kernels). + +Tebd input modes +---------------- +Concat folds the type embedding into the layer-1 input; the operator receives +an empty ``gate_table``. Strip runs the embedding on the radial channel alone +and multiplies the output by the type-pair gate +``1 + embeddings_strip(tebd_pair) (* sw when smooth_type_embedding)``; the +wrapper precomputes the gate table (``(T or T^2, ng)``) from the strip +network, so the operator never evaluates it per edge. + +Usage and pitfalls +------------------ +* The forward returns ``(grrg, rot_mat)`` plus five auxiliary tensors saved + for the backward (moment, CSR edge order, folded layer-1 pair table, and + the transposed ``pre2`` / ``g`` activation spills). The auxiliaries never + receive gradients; ``set_materialize_grads(False)`` skips their zero fill. +* The activation spills cost ``(2 * neuron[1] + neuron[2]) * 4`` bytes per + edge (768 B/edge for the (32, 64, 128) net). This is far below the autograd + tape of the reference path, but at tens of millions of edges it is + gigabytes; the buffers live only between forward and backward. +* Weight tensors must be fp32 and contiguous; ``edge_vec`` may be fp32 or + fp64 (the backward returns the gradient in the leaf dtype and the internal + compute is fp32 either way). The caller-facing gate lives in + :meth:`DescrptDPA1._fused_eligible`. +* Constants captured by the op call (weights, statistics slices, the strip + gate table) are prepared per call rather than cached on the module: values + produced while a fake / proxy tensor mode is active must never leak into + eager state. +""" + +from typing import ( + Any, +) + +import torch + +from deepmd.kernels.triton.dpa1.activation import ( + ACT_CODES, +) + +__all__ = [ + "dpa1_graph_descriptor", + "ensure_registered", + "op_available", +] + +_TILE_EDGES = 128 # matches kTileEdges in dpa1_graph_descriptor.cu + + +def op_available() -> bool: + """Whether the C++ ``deepmd::dpa1_graph_descriptor`` op is loaded.""" + op = getattr(torch.ops.deepmd, "dpa1_graph_descriptor", None) + return isinstance(op, torch._ops.OpOverloadPacket) + + +# ====================================================================== +# Fake (meta) implementations +# ====================================================================== +def _forward_fake( + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + atype: torch.Tensor, + type_embedding: torch.Tensor, + davg: torch.Tensor, + dstd: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + idt1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, + idt2: torch.Tensor, + w3: torch.Tensor, + b3: torch.Tensor, + idt3: torch.Tensor, + gate_table: torch.Tensor, + act: int, + type_one_side: int, + concat_tebd: int, + write_rotation: int, + smooth: int, + axis: int, + resnet2: int, + resnet3: int, + rcut: float, + rcut_smth: float, + protection: float, + nnei: float, +) -> tuple[torch.Tensor, ...]: + n_edge = edge_vec.shape[0] + n_node = atype.shape[0] + n2 = w2.shape[1] + ng = w3.shape[1] + ntypes, tebd_dim = type_embedding.shape + out_dim = ng * axis + (tebd_dim if concat_tebd else 0) + # Strip mode carries the type pairs in the gate table; the layer-1 pair + # table degenerates to its single bias row. + strip = gate_table.shape[0] > 0 + n_pairs = 1 if strip else (ntypes if type_one_side else ntypes * ntypes) + e_pad = (n_edge + _TILE_EDGES - 1) // _TILE_EDGES * _TILE_EDGES + dev = edge_vec.device + return ( + torch.empty(n_node, out_dim, dtype=torch.float32, device=dev), + torch.empty( + n_node if write_rotation else 0, + ng, + 3, + dtype=torch.float32, + device=dev, + ), + torch.empty(n_node, 4, ng, dtype=torch.float32, device=dev), + torch.empty(n_edge, dtype=torch.int32, device=dev), + torch.empty(n_pairs, w1.shape[1], dtype=torch.float32, device=dev), + torch.empty(n2, e_pad, dtype=torch.float32, device=dev), + torch.empty(ng, e_pad, dtype=torch.float32, device=dev), + ) + + +def _backward_fake( + d_grrg: torch.Tensor, + d_rot_mat: torch.Tensor | None, + gr: torch.Tensor, + edge_order: torch.Tensor, + pair_table: torch.Tensor, + pre2_saved: torch.Tensor, + g_saved: torch.Tensor, + edge_vec: torch.Tensor, + *args: Any, +) -> torch.Tensor: + return torch.empty_like(edge_vec) + + +# ====================================================================== +# Autograd bridge +# ====================================================================== +def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None: + ( + edge_vec, + edge_index, + edge_mask, + atype, + type_embedding, + davg, + dstd, + w1, + b1, + idt1, + w2, + b2, + idt2, + w3, + b3, + idt3, + gate_table, + act, + type_one_side, + concat_tebd, + _write_rotation, + smooth, + axis, + resnet2, + resnet3, + rcut, + rcut_smth, + protection, + nnei, + ) = inputs + # gr, edge_order, pair_table, pre2_saved, g_saved; the type embedding is + # not saved -- the backward re-reads layer 1 through the folded pair table + # and the strip gate through the gate table. + aux = output[2:] + ctx.save_for_backward( + *aux, + edge_vec, + edge_index, + edge_mask, + atype, + davg, + dstd, + w1, + b1, + idt1, + w2, + b2, + idt2, + w3, + b3, + idt3, + gate_table, + ) + ctx.n_aux = len(aux) + ctx.scalars = ( + act, + type_one_side, + smooth, + axis, + resnet2, + resnet3, + rcut, + rcut_smth, + protection, + nnei, + ) + # The auxiliary outputs never receive real gradients; skipping their zero + # materialization saves several full-size fills per step. + ctx.set_materialize_grads(False) + + +def _backward( + ctx: Any, + d_grrg: torch.Tensor, + d_rot_mat: torch.Tensor | None, + *d_aux: Any, +) -> tuple: + saved = ctx.saved_tensors + aux = saved[: ctx.n_aux] + ( + edge_vec, + edge_index, + edge_mask, + atype, + davg, + dstd, + w1, + b1, + idt1, + w2, + b2, + idt2, + w3, + b3, + idt3, + gate_table, + ) = saved[ctx.n_aux :] + ( + act, + type_one_side, + smooth, + axis, + resnet2, + resnet3, + rcut, + rcut_smth, + protection, + nnei, + ) = ctx.scalars + # rot_mat is unused by the energy fitting; autograd then hands None, which + # the op schema accepts (Tensor? d_rot_mat). + d_edge_vec = torch.ops.deepmd.dpa1_graph_descriptor_backward( + d_grrg, + d_rot_mat, + *aux, + edge_vec, + edge_index, + edge_mask, + atype, + davg, + dstd, + w1, + b1, + idt1, + w2, + b2, + idt2, + w3, + b3, + idt3, + gate_table, + act, + type_one_side, + smooth, + axis, + resnet2, + resnet3, + rcut, + rcut_smth, + protection, + nnei, + ) + return (d_edge_vec,) + (None,) * 28 + + +# ====================================================================== +# CPU reference implementations +# ====================================================================== +def _act(z: torch.Tensor, act: int) -> torch.Tensor: + return torch.tanh(z) if act == 0 else torch.nn.functional.silu(z) + + +def _cpu_pair_table( + type_embedding: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + type_one_side: int, + strip: bool, +) -> torch.Tensor: + """Fold the type-embedding rows of W1 (and b1) into the per-pair table. + + Strip mode has no type term in layer 1: the table is the single bias row. + """ + n1 = w1.shape[1] + bias = b1 if b1.numel() else torch.zeros(n1, dtype=w1.dtype, device=w1.device) + if strip: + return bias.reshape(1, n1).contiguous() + ntypes, tebd_dim = type_embedding.shape + nei = type_embedding @ w1[1 : 1 + tebd_dim] # (ntypes, n1) + if type_one_side: + return (nei + bias).contiguous() + center = type_embedding @ w1[1 + tebd_dim : 1 + 2 * tebd_dim] + return ( + (center[:, None, :] + nei[None, :, :] + bias) + .reshape(ntypes * ntypes, -1) + .contiguous() + ) + + +def _cpu_layer( + x: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + idt: torch.Tensor, + act: int, + resnet: int, +) -> tuple[torch.Tensor, torch.Tensor]: + pre = x @ w + if b.numel(): + pre = pre + b + y = _act(pre, act) + if idt.numel(): + y = y * idt + if resnet: + din, dout = w.shape + if dout == din: + y = y + x + elif dout == 2 * din: + y = y + torch.cat([x, x], dim=-1) + return pre, y + + +def _cpu_embedding( + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + atype: torch.Tensor, + pair_table: torch.Tensor, + gate_table: torch.Tensor, + davg: torch.Tensor, + dstd: torch.Tensor, + w1: torch.Tensor, + idt1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, + idt2: torch.Tensor, + w3: torch.Tensor, + b3: torch.Tensor, + idt3: torch.Tensor, + act: int, + type_one_side: int, + smooth: int, + ntypes: int, + resnet2: int, + resnet3: int, + rcut: float, + rcut_smth: float, + protection: float, +) -> tuple[torch.Tensor, ...]: + """Environment matrix, embedding MLP and (strip) type-pair gate (fp32). + + Returns ``(rr, pre2, pre3, g, gg)`` with ``gg`` the gated layer-3 output + that feeds the moment (``gg is g`` in concat mode). + """ + ev = edge_vec.to(torch.float32) + src, dst = edge_index[0], edge_index[1] + center_type, nei_type = atype[dst], atype[src] + length = ev.norm(dim=-1, keepdim=True) + length = length + (~edge_mask[:, None]).to(length.dtype) + q = length + protection + u = ((length - rcut_smth) / (rcut - rcut_smth)).clamp(0.0, 1.0) + sw = u**3 * (-6 * u**2 + 15 * u - 10) + 1.0 + em = torch.cat([sw / q, ev * (sw / q**2)], dim=-1) + rr = (em - davg[center_type]) / dstd[center_type] + strip = gate_table.shape[0] > 0 + pair_idx = nei_type if type_one_side else center_type * ntypes + nei_type + pre1 = rr[:, :1] * w1[0] + pair_table[0 if strip else pair_idx] + h1 = _act(pre1, act) + if idt1.numel(): + h1 = h1 * idt1 + pre2, h2 = _cpu_layer(h1, w2, b2, idt2, act, resnet2) + pre3, g = _cpu_layer(h2, w3, b3, idt3, act, resnet3) + gg = g + if strip: + gate = gate_table[pair_idx] + if smooth: + gate = gate * sw + gg = g * (1.0 + gate) + return rr, pre2, pre3, g, gg + + +def _cpu_outputs( + rr: torch.Tensor, + gg: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + atype: torch.Tensor, + type_embedding: torch.Tensor | None, + ng: int, + axis: int, + concat_tebd: int, + nnei: float, + n_node: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ggm = gg * edge_mask[:, None].to(gg.dtype) + outer = rr[:, :, None] * ggm[:, None, :] # (E, 4, ng) + gr = torch.zeros(n_node, 4, ng, dtype=gg.dtype, device=gg.device) + gr.index_add_(0, edge_index[1], outer) + gr = gr / nnei + gr_t = gr.permute(0, 2, 1) # (N, ng, 4) + grrg = torch.matmul(gr_t, gr[:, :, :axis]).reshape(n_node, ng * axis) + rot_mat = gr_t[:, :, 1:4].contiguous() + if concat_tebd: + grrg = torch.cat([grrg, type_embedding[atype]], dim=-1) + return grrg.contiguous(), rot_mat, gr + + +def _cpu_forward( + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + atype: torch.Tensor, + type_embedding: torch.Tensor, + davg: torch.Tensor, + dstd: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + idt1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, + idt2: torch.Tensor, + w3: torch.Tensor, + b3: torch.Tensor, + idt3: torch.Tensor, + gate_table: torch.Tensor, + act: int, + type_one_side: int, + concat_tebd: int, + write_rotation: int, + smooth: int, + axis: int, + resnet2: int, + resnet3: int, + rcut: float, + rcut_smth: float, + protection: float, + nnei: float, +) -> tuple[torch.Tensor, ...]: + n_edge = edge_vec.shape[0] + n_node = atype.shape[0] + n2, ng = w2.shape[1], w3.shape[1] + ntypes = type_embedding.shape[0] + strip = gate_table.shape[0] > 0 + pair_table = _cpu_pair_table(type_embedding, w1, b1, type_one_side, strip) + rr, pre2, pre3, g, gg = _cpu_embedding( + edge_vec, + edge_index, + edge_mask, + atype, + pair_table, + gate_table, + davg, + dstd, + w1, + idt1, + w2, + b2, + idt2, + w3, + b3, + idt3, + act, + type_one_side, + smooth, + ntypes, + resnet2, + resnet3, + rcut, + rcut_smth, + protection, + ) + grrg, rot_mat, gr = _cpu_outputs( + rr, + gg, + edge_index, + edge_mask, + atype, + type_embedding, + ng, + axis, + concat_tebd, + nnei, + n_node, + ) + if not write_rotation: + rot_mat = rot_mat.new_empty(0, ng, 3) + e_pad = (n_edge + _TILE_EDGES - 1) // _TILE_EDGES * _TILE_EDGES + dev = edge_vec.device + pre2_saved = torch.zeros(n2, e_pad, dtype=torch.float32, device=dev) + pre2_saved[:, :n_edge] = pre2.t() + g_saved = torch.zeros(ng, e_pad, dtype=torch.float32, device=dev) + g_saved[:, :n_edge] = (g if act == 0 else pre3).t() + edge_order = torch.arange(n_edge, dtype=torch.int32, device=dev) + return grrg, rot_mat, gr, edge_order, pair_table, pre2_saved, g_saved + + +def _cpu_backward( + d_grrg: torch.Tensor, + d_rot_mat: torch.Tensor | None, + gr: torch.Tensor, + edge_order: torch.Tensor, + pair_table: torch.Tensor, + pre2_saved: torch.Tensor, + g_saved: torch.Tensor, + edge_vec: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + atype: torch.Tensor, + davg: torch.Tensor, + dstd: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + idt1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, + idt2: torch.Tensor, + w3: torch.Tensor, + b3: torch.Tensor, + idt3: torch.Tensor, + gate_table: torch.Tensor, + act: int, + type_one_side: int, + smooth: int, + axis: int, + resnet2: int, + resnet3: int, + rcut: float, + rcut_smth: float, + protection: float, + nnei: float, +) -> torch.Tensor: + n_node = atype.shape[0] + ng = w3.shape[1] + strip = gate_table.shape[0] > 0 + n_pairs = gate_table.shape[0] if strip else pair_table.shape[0] + ntypes = n_pairs if type_one_side else round(n_pairs**0.5) + ev = edge_vec.detach().clone().requires_grad_(True) + with torch.enable_grad(): + rr, _pre2, _pre3, _g, gg = _cpu_embedding( + ev, + edge_index, + edge_mask, + atype, + pair_table, + gate_table, + davg, + dstd, + w1, + idt1, + w2, + b2, + idt2, + w3, + b3, + idt3, + act, + type_one_side, + smooth, + ntypes, + resnet2, + resnet3, + rcut, + rcut_smth, + protection, + ) + grrg, rot_mat, _gr = _cpu_outputs( + rr, + gg, + edge_index, + edge_mask, + atype, + None, + ng, + axis, + 0, + nnei, + n_node, + ) + loss = (grrg * d_grrg[:, : grrg.shape[1]].to(grrg.dtype)).sum() + if d_rot_mat is not None: + loss = loss + (rot_mat * d_rot_mat.to(rot_mat.dtype)).sum() + (d_ev,) = torch.autograd.grad(loss, ev) + return d_ev.to(edge_vec.dtype) + + +# ====================================================================== +# Registration and the public wrapper +# ====================================================================== +_cpu_library: torch.library.Library | None = None + + +def ensure_registered() -> None: + """Register the fake / backward / CPU implementations for the ops. + + Idempotent; a no-op when the C++ operator library is not loaded (so the + plain pt inference path, which never dispatches here, does not require + the ops). + """ + global _cpu_library + if _cpu_library is not None or not op_available(): + return + torch.library.register_fake("deepmd::dpa1_graph_descriptor")(_forward_fake) + torch.library.register_fake("deepmd::dpa1_graph_descriptor_backward")( + _backward_fake + ) + torch.library.register_autograd( + "deepmd::dpa1_graph_descriptor", _backward, setup_context=_setup_context + ) + _cpu_library = torch.library.Library("deepmd", "IMPL") + _cpu_library.impl("dpa1_graph_descriptor", _cpu_forward, "CPU") + _cpu_library.impl("dpa1_graph_descriptor_backward", _cpu_backward, "CPU") + + +def _strip_gate_table(se: Any, type_embedding: torch.Tensor) -> torch.Tensor: + """Type-pair gate table (T or T^2, ng) through the strip embedding net. + + One-side embeds the neighbor-type rows; two-side embeds every + ``[neighbor, center]`` pair, matching the dense reference's + ``cal_g_strip`` input layout. + """ + if se.type_one_side: + pairs = type_embedding + else: + ntypes, tebd_dim = type_embedding.shape + nei = type_embedding.view(1, ntypes, tebd_dim).expand(ntypes, ntypes, tebd_dim) + center = type_embedding.view(ntypes, 1, tebd_dim).expand( + ntypes, ntypes, tebd_dim + ) + pairs = torch.cat([nei, center], dim=-1).reshape(ntypes * ntypes, 2 * tebd_dim) + return se.embeddings_strip[0].call(pairs).to(torch.float32).contiguous() + + +def dpa1_graph_descriptor( + desc: Any, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor, + write_rotation: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """Evaluate the fused DPA1 descriptor from the edge stream. + + Numerically equivalent to the dpmodel reference on the attention-free + configuration (up to fp32 summation order), in both tebd input modes: + concat, and strip with or without the smooth type-embedding gate. + ``edge_vec`` may arrive fp64 (the model-agnostic ``.pt2`` ABI); the cast + to the fp32 compute happens inside the operator, whose backward returns + ``d_edge_vec`` in the leaf dtype. + + Parameters + ---------- + desc : DescrptDPA1 + The pt_expt descriptor module; must satisfy + ``desc._fused_eligible("cuda")``. + graph : NeighborGraph + The lowered neighbor graph: ``edge_vec`` (E, 3) is the autograd leaf + for the force / virial backward, ``edge_index`` (2, E) the + ``[src, dst]`` endpoints, ``edge_mask`` (E,) the valid-edge mask. + atype : torch.Tensor + Flat node atom types with shape (N,), int64. + type_embedding : torch.Tensor + Type embedding table with shape (ntypes + 1, tebd_dim), fp32. + write_rotation : bool, default: True + Whether to materialize the equivariant rotation output. Energy-only + inference can disable it. + + Returns + ------- + grrg : torch.Tensor + Descriptor with shape (N, ng * axis [+ tebd_dim]), fp32. + rot_mat : torch.Tensor + Equivariant rotation matrix with shape (N, ng, 3), fp32. + """ + ensure_registered() + se = desc.se_atten + layers = se.embeddings[0].layers + empty = layers[0].w.new_empty(0) + + def optional(t: torch.Tensor | None) -> torch.Tensor: + return t.contiguous() if t is not None else empty + + if se.tebd_input_mode == "strip": + gate_table = _strip_gate_table(se, type_embedding) + smooth = int(se.smooth) + else: + gate_table = empty.reshape(0, 0) + smooth = 0 + w1, w2, w3 = (layer.w.contiguous() for layer in layers) + grrg, rot_mat, *_aux = torch.ops.deepmd.dpa1_graph_descriptor( + graph.edge_vec.contiguous(), + graph.edge_index.contiguous(), + graph.edge_mask.contiguous(), + atype.contiguous(), + type_embedding.contiguous(), + # mean / stddev are slot-independent; slot 0 is the canonical (T, 4). + se.mean[:, 0, :].contiguous(), + se.stddev[:, 0, :].contiguous(), + w1, + optional(layers[0].b), + optional(layers[0].idt), + w2, + optional(layers[1].b), + optional(layers[1].idt), + w3, + optional(layers[2].b), + optional(layers[2].idt), + gate_table, + ACT_CODES[str(layers[0].activation_function).lower()], + int(se.type_one_side), + int(desc.concat_output_tebd), + int(write_rotation), + smooth, + int(se.axis_neuron), + int(layers[1].resnet), + int(layers[2].resnet), + float(se.rcut), + float(se.rcut_smth), + float(se.env_protection), + float(se.nnei), + ) + return grrg, rot_mat diff --git a/deepmd/kernels/cuda/dpa1/graph_energy_force.py b/deepmd/kernels/cuda/dpa1/graph_energy_force.py new file mode 100644 index 0000000000..2cc262290e --- /dev/null +++ b/deepmd/kernels/cuda/dpa1/graph_energy_force.py @@ -0,0 +1,430 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the end-to-end DPA1 graph-lower energy-force operator. + +The CUDA operator ``deepmd::dpa1_graph_energy_force`` (see +``source/op/pt/dpa1_graph_energy_force.cu``) fuses the whole DPA1 (``se_atten``, +attention-free) energy path -- descriptor, energy fitting, and the analytic +force / virial assembly -- into one graph node. It drives the descriptor and +fitting mega kernels and their backwards directly, computing the force from the +reduced energy internally (``dE_redu/d(atom_energy) == 1``), so no autograd tape +is built and the force is returned as a value. + +It is numerically identical to the separate-operator path (descriptor + fitting +forwards with the force from ``autograd.grad``): the same operator backwards run +the same fp32 arithmetic in the same order. The fusion removes the autograd +machinery, the output-agnostic per-component grad loop, and the inter-operator +array glue. It is selected at freeze time by ``DP_CUDA_INFER >= 2``; the +separate-operator level-1 path is unchanged. + +The operator is forward-only (the force is an output, not a gradient), so only +the fake and CPU implementations are registered -- no autograd bridge. +""" + +from typing import ( + Any, +) + +import torch + +from deepmd.kernels.cuda.dpa1.graph_descriptor import ( + _strip_gate_table, +) +from deepmd.kernels.cuda.dpa1.graph_descriptor import ( + ensure_registered as ensure_descriptor_registered, +) +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, +) +from deepmd.kernels.triton.dpa1.activation import ( + ACT_CODES, +) + +__all__ = [ + "dpa1_graph_energy_force", + "ensure_registered", + "op_available", +] + + +def op_available() -> bool: + """Whether the C++ ``deepmd::dpa1_graph_energy_force`` op is loaded.""" + op = getattr(torch.ops.deepmd, "dpa1_graph_energy_force", None) + return isinstance(op, torch._ops.OpOverloadPacket) + + +def _fake( + 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, + atype: torch.Tensor, + n_node: torch.Tensor, + ownership: torch.Tensor, + type_embedding: torch.Tensor, + davg: torch.Tensor, + dstd: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + idt1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, + idt2: torch.Tensor, + w3: torch.Tensor, + b3: torch.Tensor, + idt3: torch.Tensor, + gate_table: torch.Tensor, + act: int, + type_one_side: int, + concat_tebd: int, + smooth: int, + axis: int, + resnet2: int, + resnet3: int, + rcut: float, + rcut_smth: float, + protection: float, + nnei: float, + 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, + bias_atom_e: torch.Tensor, + fit_act: int, + node_capacity: int, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, ...]: + nf = n_node.shape[0] + n = atype.shape[0] + dev = edge_vec.device + # Force / virial assemble in the model compute precision (descriptor weight + # dtype), matching the real operator; the energy reduction stays fp64. + fprec = w1.dtype + return ( + torch.empty(nf, 1, dtype=torch.float64, device=dev), + torch.empty(n, 1, dtype=torch.float64, device=dev), + torch.empty(n, 3, dtype=fprec, device=dev), + torch.empty(nf, 3, 3, dtype=fprec, device=dev), + torch.empty(n if do_atomic_virial else 0, 3, 3, dtype=fprec, device=dev), + ) + + +def _cpu( + 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, + atype: torch.Tensor, + n_node: torch.Tensor, + ownership: torch.Tensor, + type_embedding: torch.Tensor, + davg: torch.Tensor, + dstd: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + idt1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, + idt2: torch.Tensor, + w3: torch.Tensor, + b3: torch.Tensor, + idt3: torch.Tensor, + gate_table: torch.Tensor, + act: int, + type_one_side: int, + concat_tebd: int, + smooth: int, + axis: int, + resnet2: int, + resnet3: int, + rcut: float, + rcut_smth: float, + protection: float, + nnei: float, + 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, + bias_atom_e: torch.Tensor, + fit_act: int, + node_capacity: int, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, ...]: + """Trace-time reference: the same sequence over the sub-operators' CPU + implementations, so the exported graph's sample evaluation is faithful. + """ + grrg, _rot, gr, order, pair_table, pre2_saved, g_saved = ( + torch.ops.deepmd.dpa1_graph_descriptor( + edge_vec, + edge_index, + edge_mask, + atype, + type_embedding, + davg, + dstd, + w1, + b1, + idt1, + w2, + b2, + idt2, + w3, + b3, + idt3, + gate_table, + act, + type_one_side, + concat_tebd, + 0, + smooth, + axis, + resnet2, + resnet3, + rcut, + rcut_smth, + protection, + nnei, + ) + ) + atom_e_raw, fit_saved = torch.ops.deepmd.graph_fitting( + grrg, + atype, + fit_ws, + fit_bs, + fit_idts, + fit_resnets, + w_head, + b_head, + bias_atom_e, + fit_act, + ) + owned = ownership[:, None].to(atom_e_raw.dtype) + energy_seed = owned + atom_e = atom_e_raw * owned + nf = n_node.shape[0] + frame_id = torch.repeat_interleave( + torch.arange(nf, dtype=torch.long, device=n_node.device), n_node + ) + 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 + ) + del grrg, fit_saved + g_e = torch.ops.deepmd.dpa1_graph_descriptor_backward( + d_grrg, + None, + gr, + order, + pair_table, + pre2_saved, + g_saved, + edge_vec, + edge_index, + edge_mask, + atype, + davg, + dstd, + w1, + b1, + idt1, + w2, + b2, + idt2, + w3, + b3, + idt3, + gate_table, + act, + type_one_side, + smooth, + axis, + resnet2, + resnet3, + rcut, + rcut_smth, + protection, + nnei, + ) + # Assemble force / virial in the descriptor weight precision, matching + # ``_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( + g_e.to(fprec), + edge_vec.to(fprec), + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + n_node, + node_capacity, + do_atomic_virial, + ) + return energy, atom_e, force, virial, atom_virial + + +_cpu_library: torch.library.Library | None = None + + +def ensure_registered() -> None: + """Register the fake and CPU implementations for the op. + + Idempotent; a no-op when the C++ operator library is not loaded. The + sub-operators must be registered too, since the CPU reference dispatches + through them. + """ + global _cpu_library + if _cpu_library is not None or not op_available(): + return + ensure_descriptor_registered() + ensure_fitting_registered() + ensure_force_registered() + torch.library.register_fake("deepmd::dpa1_graph_energy_force")(_fake) + _cpu_library = torch.library.Library("deepmd", "IMPL") + _cpu_library.impl("dpa1_graph_energy_force", _cpu, "CPU") + + +def dpa1_graph_energy_force( + desc: Any, + fit: Any, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + node_capacity: int, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused energy / force / virial from the edge stream. + + Parameters + ---------- + desc : DescrptDPA1 + The pt_expt descriptor module; must satisfy ``desc._fused_eligible("cuda")``. + fit : EnergyFittingNet + The pt_expt fitting module; must satisfy + :func:`~deepmd.kernels.cuda.graph_fitting.fitting_eligible`. + graph : NeighborGraph + The lowered neighbor graph (``edge_vec``, ``edge_index``, ``edge_mask``, + ``n_node``) with destination/source CSR permutations. + atype : torch.Tensor + Flat node atom types with shape (N,), int64. + type_embedding : torch.Tensor + Type embedding table with shape (ntypes + 1, tebd_dim), fp32. + ownership : torch.Tensor + Energy-contributing node mask with shape (N,), bool. + atom_bias : torch.Tensor + Combined fitting and atomic-model bias with shape (ntypes,). + node_capacity : int + Padded node-axis size ``N`` (the force / atom-virial scatter bound). + do_atomic_virial : bool + Whether to also assemble the per-atom virial. + + Returns + ------- + energy : torch.Tensor + Per-frame energy with shape (nf, 1), fp64. + atom_energy : torch.Tensor + Per-atom energy with shape (N, 1), fp64. + force : torch.Tensor + Per-atom force with shape (N, 3), ``edge_vec`` dtype. + virial : torch.Tensor + Per-frame virial with shape (nf, 3, 3), ``edge_vec`` dtype. + atom_virial : torch.Tensor + Per-atom virial with shape (N, 3, 3) when requested, else empty (0, 3, 3). + """ + ensure_registered() + se = desc.se_atten + layers = se.embeddings[0].layers + empty = layers[0].w.new_empty(0) + + def optional(t: torch.Tensor | None) -> torch.Tensor: + return t.contiguous() if t is not None else empty + + if se.tebd_input_mode == "strip": + gate_table = _strip_gate_table(se, type_embedding) + smooth = int(se.smooth) + else: + gate_table = empty.reshape(0, 0) + 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) + 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( + "DPA1 fused inference requires destination/source CSR topology" + ) + return torch.ops.deepmd.dpa1_graph_energy_force( + graph.edge_vec.contiguous(), + graph.edge_index.contiguous(), + graph.edge_mask.contiguous(), + graph.destination_order.contiguous(), + graph.destination_row_ptr.contiguous(), + graph.source_order.contiguous(), + graph.source_row_ptr.contiguous(), + atype.contiguous(), + graph.n_node, + ownership.contiguous(), + type_embedding.contiguous(), + se.mean[:, 0, :].contiguous(), + se.stddev[:, 0, :].contiguous(), + w1, + optional(layers[0].b), + optional(layers[0].idt), + w2, + optional(layers[1].b), + optional(layers[1].idt), + w3, + optional(layers[2].b), + optional(layers[2].idt), + gate_table, + ACT_CODES[str(layers[0].activation_function).lower()], + int(se.type_one_side), + int(desc.concat_output_tebd), + smooth, + int(se.axis_neuron), + int(layers[1].resnet), + int(layers[2].resnet), + float(se.rcut), + float(se.rcut_smth), + float(se.env_protection), + float(se.nnei), + [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 + ), + atom_bias.to(torch.float64).contiguous(), + ACT_CODES[str(hidden[0].activation_function).lower()], + node_capacity, + do_atomic_virial, + ) diff --git a/deepmd/kernels/cuda/edge_force_virial.py b/deepmd/kernels/cuda/edge_force_virial.py new file mode 100644 index 0000000000..22cbe58007 --- /dev/null +++ b/deepmd/kernels/cuda/edge_force_virial.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the fused force / virial assembly of the graph lower. + +The CUDA operator ``deepmd::edge_force_virial`` (see +``source/op/pt/edge_force_virial.cu``) scatters the per-edge energy gradient +``g_e = dE/d(edge_vec)`` into per-atom force, per-atom virial (optional) and +per-frame virial through destination/source CSR reductions, replacing the array-API +:func:`~deepmd.dpmodel.utils.neighbor_graph.edge_force_virial` chain of +``index_add`` / outer-product / ``segment_sum`` kernels. It is +descriptor-agnostic: any graph-lowered model whose force path differentiates +the energy w.r.t. ``edge_vec`` can dispatch here. + +Usage and pitfalls +------------------ +* ``node_capacity`` is declared ``SymInt`` in the op schema. With a plain + ``int`` the traced graph would specialize the padded node count to the + trace-time value and a resized deployment would read out of bounds; the + ``SymInt`` keeps it symbolic through ``make_fx`` / ``torch.export``. +* The op is dispatched downstream of a ``torch.autograd.grad`` call, so it + needs no registered backward of its own; only the fake and CPU + implementations are required for tracing. +* When the atomic virial is not requested the op returns an empty + ``(0, 3, 3)`` tensor instead of skipping the output (the schema is static); + the caller maps it to ``None``. +* CSR rows describe topology rather than validity. Masked entries may remain + inside a row, so the kernel applies ``edge_mask`` to both incidence + reductions. +""" + +import torch + +__all__ = [ + "canonical_edge_force_virial", + "canonical_op_available", + "edge_force_virial", + "ensure_registered", + "op_available", +] + + +def op_available() -> bool: + """Whether the C++ ``deepmd::edge_force_virial`` op is loaded.""" + op = getattr(torch.ops.deepmd, "edge_force_virial", None) + return isinstance(op, torch._ops.OpOverloadPacket) + + +def canonical_op_available() -> bool: + """Whether the compact canonical force operator is loaded.""" + op = getattr(torch.ops.deepmd, "canonical_edge_force_virial", None) + return isinstance(op, torch._ops.OpOverloadPacket) + + +def _fake( + 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, + node_capacity: int, + want_atom_virial: bool, +) -> tuple[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), + ) + + +def _canonical_fake( + g_e: 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, + node_capacity: int, + want_atom_virial: bool, +) -> tuple[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), + ) + + +def _cpu( + 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, + node_capacity: int, + want_atom_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from deepmd.dpmodel.utils.neighbor_graph import edge_force_virial as reference + + force, atom_virial, virial = reference( + g_e, + edge_vec, + edge_index, + edge_mask, + n_node_per_frame, + node_capacity=node_capacity, + ) + if not want_atom_virial: + atom_virial = atom_virial.new_zeros(0, 3, 3) + return force, atom_virial, virial + + +def _canonical_cpu( + g_e: 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, + node_capacity: int, + want_atom_virial: bool, +) -> tuple[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( + torch.arange(node_count, dtype=torch.int64, device=edge_vec.device), + destination_row_ptr[1:] - destination_row_ptr[:-1], + output_size=physical_edge_count, + ) + source_by_row = torch.repeat_interleave( + torch.arange(node_count, dtype=torch.int64, device=edge_vec.device), + source_row_ptr[1:] - source_row_ptr[:-1], + output_size=physical_edge_count, + ) + source = torch.zeros( + edge_vec.shape[0], + dtype=torch.int64, + device=edge_vec.device, + ) + source[source_order[:physical_edge_count].to(torch.int64)] = source_by_row + destination_storage = torch.zeros_like(source) + destination_storage[:physical_edge_count] = destination + edge_index = torch.stack((source, destination_storage)) + edge_mask = ( + torch.arange( + edge_vec.shape[0], + dtype=torch.int64, + device=edge_vec.device, + ) + < physical_edge_count + ) + return _cpu( + g_e, + edge_vec, + edge_index, + edge_mask, + torch.arange( + edge_vec.shape[0], + dtype=source_order.dtype, + device=edge_vec.device, + ), + destination_row_ptr, + source_order, + source_row_ptr, + n_node_per_frame, + node_capacity, + want_atom_virial, + ) + + +_cpu_library: torch.library.Library | None = None + + +def ensure_registered() -> None: + """Register the fake and CPU implementations for the op. + + Idempotent; a no-op when the C++ operator library is not loaded. + """ + global _cpu_library + if _cpu_library is not None or not op_available(): + return + torch.library.register_fake("deepmd::edge_force_virial")(_fake) + if canonical_op_available(): + torch.library.register_fake("deepmd::canonical_edge_force_virial")( + _canonical_fake + ) + _cpu_library = torch.library.Library("deepmd", "IMPL") + _cpu_library.impl("edge_force_virial", _cpu, "CPU") + if canonical_op_available(): + _cpu_library.impl( + "canonical_edge_force_virial", + _canonical_cpu, + "CPU", + ) + + +def edge_force_virial( + 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, + node_capacity: int, + want_atom_virial: bool, +) -> tuple[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: + ``force[k] = sum_{dst=k} g_e - sum_{src=k} g_e``, the atom virial is + attributed full-to-src as ``-g_e (x) edge_vec``, and the per-frame virial + reduces the (implicit) atom virial over each frame's atoms. + + Parameters + ---------- + g_e : torch.Tensor + Per-edge energy gradient with shape (E, 3). + edge_vec : torch.Tensor + Per-edge displacement with shape (E, 3), same dtype as ``g_e``. + edge_index : torch.Tensor + ``[src, dst]`` node endpoints with shape (2, E), int64. + edge_mask : torch.Tensor + Valid-edge mask with shape (E,), bool. + destination_order, source_order : torch.Tensor + Edge permutations grouped by destination/source with shape (E,), int32 + or int64. + destination_row_ptr, source_row_ptr : torch.Tensor + Destination/source CSR offsets with shape (N + 1,), int64. + n_node_per_frame : torch.Tensor + Per-frame node counts with shape (nf,), int64. + node_capacity : int + Padded node-axis size ``N`` (may be a ``SymInt`` under tracing). + want_atom_virial : bool + Whether to materialize the per-atom virial. + + Returns + ------- + force : torch.Tensor + Per-atom force with shape (N, 3). + atom_virial : torch.Tensor + Per-atom virial with shape (N, 3, 3), or an empty (0, 3, 3) tensor + when not requested. + virial : torch.Tensor + Per-frame virial with shape (nf, 3, 3). + """ + ensure_registered() + return torch.ops.deepmd.edge_force_virial( + g_e.contiguous(), + edge_vec.contiguous(), + edge_index.contiguous(), + edge_mask.contiguous(), + destination_order.contiguous(), + destination_row_ptr.contiguous(), + source_order.contiguous(), + source_row_ptr.contiguous(), + n_node_per_frame, + node_capacity, + want_atom_virial, + ) + + +def canonical_edge_force_virial( + g_e: 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, + node_capacity: int, + want_atom_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Assemble force and virial from a compact canonical edge stream. + + Parameters + ---------- + g_e + Per-edge energy gradient with shape ``(S, 3)``. + edge_vec + Per-edge displacement with shape ``(S, 3)``. + destination_row_ptr, source_row_ptr + Destination and source CSR offsets with shape ``(N + 1,)``. + source_order + Edge storage positions grouped by source with shape ``(S,)``. + n_node_per_frame + Per-frame node counts with shape ``(nf,)``. + node_capacity + Flat node count ``N``. + want_atom_virial + Whether to materialize the per-node virial. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor, torch.Tensor] + Force, optional atom virial, and frame virial. + """ + ensure_registered() + return torch.ops.deepmd.canonical_edge_force_virial( + g_e.contiguous(), + edge_vec.contiguous(), + destination_row_ptr.contiguous(), + source_row_ptr.contiguous(), + source_order.contiguous(), + n_node_per_frame, + node_capacity, + want_atom_virial, + ) diff --git a/deepmd/kernels/cuda/graph_fitting.py b/deepmd/kernels/cuda/graph_fitting.py new file mode 100644 index 0000000000..fb8b7624c5 --- /dev/null +++ b/deepmd/kernels/cuda/graph_fitting.py @@ -0,0 +1,309 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the fused energy fitting operator of the graph lower. + +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. + +The operator is descriptor-agnostic: any graph-lowered energy model whose +fitting is a plain MLP over the flat node axis (see +:func:`fitting_eligible`) dispatches here, regardless of which descriptor +produced the input. + +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 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. +""" + +from typing import ( + Any, +) + +import torch + +from deepmd.kernels.triton.dpa1.activation import ( + ACT_CODES, +) + +__all__ = [ + "ensure_registered", + "fitting_eligible", + "graph_fitting", + "op_available", +] + + +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) + + +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. + + Parameters + ---------- + fit : EnergyFittingNet + The pt_expt fitting module. + + Returns + ------- + bool + ``True`` when the fused operator reproduces the reference forward. + """ + if fit.numb_fparam or fit.numb_aparam or fit.dim_case_embd: + return False + if getattr(fit, "exclude_types", None): + return False + if not fit.mixed_types or len(fit.nets._networks) != 1: + return False + layers = fit.nets[0].layers + if len(layers) < 2: + return False + *hidden, head = layers + acts = {str(layer.activation_function).lower() for layer in hidden} + if len(acts) != 1 or acts.pop() not in ACT_CODES: + return False + if str(head.activation_function).lower() not in ("none", "linear"): + return False + if head.w.shape[1] != 1 or head.idt is not None or head.resnet: + return False + tensors = [ + tensor + for layer in layers + for tensor in (layer.w, layer.b, layer.idt) + if tensor is not None + ] + if any(tensor.dtype != torch.float32 for tensor in tensors): + return False + if any( + 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) + + +# ====================================================================== +# Fake (meta) implementations and the autograd bridge +# ====================================================================== +def _forward_fake( + 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]: + n_node = x.shape[0] + total_width = sum(int(w.shape[1]) for w in ws) + return ( + x.new_empty(n_node, 1, dtype=torch.float64), + x.new_empty(n_node * total_width, dtype=torch.float32), + ) + + +def _backward_fake( + d_e: torch.Tensor, + saved: torch.Tensor, + ws: list[torch.Tensor], + resnets: list[int], + w_head: torch.Tensor, +) -> torch.Tensor: + total_width = sum(int(w.shape[1]) for w in ws) + n_node = saved.shape[0] // total_width + return saved.new_empty(n_node, ws[0].shape[0]) + + +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 + _e, saved = output + ctx.save_for_backward(saved, w_head, *ws) + ctx.n_layers = len(ws) + ctx.resnets = resnets + ctx.set_materialize_grads(False) + + +def _backward(ctx: Any, d_e: torch.Tensor, d_saved: Any) -> tuple: + saved, w_head, *ws = ctx.saved_tensors + d_x = torch.ops.deepmd.graph_fitting_backward( + d_e, saved, list(ws), ctx.resnets, w_head + ) + none_list = [None] * ctx.n_layers + return (d_x, None, none_list, none_list, none_list, None, None, None, None, None) + + +# ====================================================================== +# CPU reference implementations +# ====================================================================== +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 = [] + cur = x.to(torch.float32) + for w, b, idt, res in zip(ws, bs, idts, resnets): + pre = cur @ w + 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]) + return e, saved + + +def _cpu_backward( + d_e: torch.Tensor, + saved: torch.Tensor, + ws: list[torch.Tensor], + resnets: list[int], + w_head: torch.Tensor, +) -> 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]: + dx = dx + dh + dh = dx + return dh + + +# ====================================================================== +# Registration and the public wrapper +# ====================================================================== +_cpu_library: torch.library.Library | None = None + + +def ensure_registered() -> None: + """Register the fake / backward / CPU implementations for the ops. + + Idempotent; a no-op when the C++ operator library is not loaded. + """ + global _cpu_library + if _cpu_library is not None or not op_available(): + return + torch.library.register_fake("deepmd::graph_fitting")(_forward_fake) + torch.library.register_fake("deepmd::graph_fitting_backward")(_backward_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") + + +def graph_fitting( + fit: Any, + descriptor: torch.Tensor, + atype: torch.Tensor, +) -> dict[str, torch.Tensor]: + """Fused energy fitting on the flat node axis. + + Drop-in for ``GeneralFitting.call_graph`` on the configuration accepted + by :func:`fitting_eligible`; the descriptor gradient flows through the + registered backward so the force / virial assembly differentiates + end-to-end. + + Parameters + ---------- + fit : EnergyFittingNet + The pt_expt fitting module. + descriptor : torch.Tensor + Flat descriptor with shape (N, nd). + atype : torch.Tensor + Flat node atom types with shape (N,), int64. + + Returns + ------- + dict[str, torch.Tensor] + ``{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) + 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 + ), + fit.bias_atom_e.to(torch.float64).reshape(-1, 1)[:, 0].contiguous(), + ACT_CODES[str(hidden[0].activation_function).lower()], + ) + return {fit.var_name: e} diff --git a/deepmd/kernels/triton/dpa1/__init__.py b/deepmd/kernels/triton/dpa1/__init__.py new file mode 100644 index 0000000000..439cfc7d8c --- /dev/null +++ b/deepmd/kernels/triton/dpa1/__init__.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Hardware-accelerated DPA1 (``se_atten``) operators. + +This package hosts ``make_fx``-composable Triton implementations of the DPA1 +descriptor hot path: the strip / dense environment convolution ``se_conv`` +(node-parallel) and the concat / graph environment convolution ``edge_conv`` +(edge-parallel). Both are internal to the ``se_atten`` descriptor; the +package-level API exposes the fused operators and the shared availability flag. +""" + +from .activation import ( + TRITON_AVAILABLE, +) +from .edge_conv import ( + edge_conv, +) +from .se_conv import ( + se_conv, +) + +__all__ = [ + "TRITON_AVAILABLE", + "edge_conv", + "se_conv", +] diff --git a/deepmd/kernels/triton/dpa1/activation.py b/deepmd/kernels/triton/dpa1/activation.py new file mode 100644 index 0000000000..b705c643bf --- /dev/null +++ b/deepmd/kernels/triton/dpa1/activation.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +# pyright: reportMissingImports=false +# ruff: noqa: ANN001, ANN201 +"""Shared last-layer activation primitives for the DPA1 (``se_atten``) Triton +inference kernels. + +Both fused environment convolutions -- the strip / dense ``se_conv`` +(node-parallel) and the concat / graph ``edge_conv`` (edge-parallel) -- inline +the embedding net's last-layer activation for value and derivative. The +activation-code map, the ``triton`` availability flag and the two ``@triton.jit`` +helpers live here so neither kernel module depends on the other. +""" + +from __future__ import ( + annotations, +) + +# Last-layer activations the fused kernels inline, mapped to the ``ACT`` code +# consumed by the Triton bodies (forward value and backward derivative). Only +# these activations are eligible for a fused path; any other keeps the dense +# reference. Both are smooth (differentiable forces), unlike relu / relu6. +ACT_CODES: dict[str, int] = {"tanh": 0, "silu": 1} + +try: + import triton + import triton.language as tl + from triton.language.extra import ( + libdevice, + ) + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover - exercised only without triton + TRITON_AVAILABLE = False + + +if TRITON_AVAILABLE: + + @triton.jit + def activation(z, ACT: tl.constexpr): + """Last-layer activation value; ``ACT`` selects ``tanh`` (0) or ``silu`` (1).""" + if ACT == 0: + a = libdevice.tanh(z) + else: + a = z * tl.sigmoid(z) + return a + + @triton.jit + def activation_grad(z, ACT: tl.constexpr): + """Activation value and its derivative for ``ACT`` in {``tanh``, ``silu``}. + + For ``tanh`` the derivative is ``1 - tanh(z)^2``; for ``silu`` it is + ``sigmoid(z) * (1 + z * (1 - sigmoid(z)))``. + """ + if ACT == 0: + a = libdevice.tanh(z) + ad = 1.0 - a * a + else: + s = tl.sigmoid(z) + a = z * s + ad = s * (1.0 + z * (1.0 - s)) + return a, ad diff --git a/deepmd/kernels/triton/dpa1/edge_conv.py b/deepmd/kernels/triton/dpa1/edge_conv.py new file mode 100644 index 0000000000..ab028eeb9c --- /dev/null +++ b/deepmd/kernels/triton/dpa1/edge_conv.py @@ -0,0 +1,900 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +# pyright: reportMissingImports=false +# ruff: noqa: ANN001, ANN202 +"""Fused graph-native environment convolution for the DPA1 (``se_atten``) +descriptor, in either concat or strip tebd-input mode. + +The graph lower represents the neighbor list as a flat edge stream of ``E`` +edges (``edge_index = [src, dst]``, ``edge_vec = r_src - r_dst``) rather than a +padded ``(node, nnei)`` grid, and replaces the neighbor-axis reduction with a +``segment_sum`` over edge centers (``dst``). For the attention-free +(``attn_layer == 0``) path the per-edge descriptor tail is + + ``h2 = act(z2) * idt + resnet(h1)`` (last embedding layer + resnet) + ``gg = h2 * (1 + tt[idx] * sw)`` (strip type-pair gate; concat: gg = h2) + ``gg = gg * edge_mask`` (drop padding edges) + ``gr[dst, k, c] += rr[k] * gg[c]`` (outer product + segment_sum) + +where ``act`` is the last layer's activation (``tanh`` or ``silu``), ``z2`` the +last embedding pre-activation, ``h1`` the penultimate activation feeding the +resnet, ``idt`` the per-channel timestep, and ``rr`` the ``(s, s*x/r, s*y/r, +s*z/r)`` per-edge environment matrix. The two tebd-input modes differ only in +the gate (``gated``): concat carries the type feature through the embedding +input (no gate); strip factorizes it into the type-pair table ``tt`` gathered +per edge by ``idx`` and scaled by the per-edge switch ``sw``, matching the dense +``se_conv`` gate. The ``1 / nnei`` normalization of ``gr`` is applied by the +descriptor after this operator, so the operator returns the unnormalized moment. + +Layout and strategy +------------------- +- **Node-parallel forward via a CSR segment reduction.** The operator builds a + destination-sorted topology internally (``argsort`` on ``dst`` plus + ``searchsorted`` for the per-node segment offsets -- integer, gradient-free), + so the forward is one program per node: it streams the node's edges (a + contiguous segment of the sorted ``order``) in ``BLOCK_E``-edge blocks, forms + ``gg`` in registers, accumulates the four moment rows and writes ``gr`` (shape + ``(N, 4, ng)``) exactly once. This is contention-free -- unlike a per-edge + ``atomic_add`` scatter, whose ~nnei-way ``dst`` collisions serialize at + production neighbor densities -- so it matches the dense ``se_conv`` register + reduction while streaming only real edges (no ``sel`` padding). ``edge_mask`` + must gate ``gg``: ``edge_env_mat`` normalizes by the per-type mean, so padding + edges carry a *nonzero* ``rr`` and would otherwise contaminate the reduction. +- **Gather backward, edge-parallel.** The backward reads the upstream moment + gradient once per edge from ``grad_gr[dst]`` (a gather) and writes ``dz2`` / + ``dh1`` / ``drr`` per edge, so it needs neither the sorted topology nor + atomics. +- **Arbitrary channel widths via padding.** ``tl.arange`` requires a + power-of-two bound; the channel axis is carried at ``next_pow2(ng)`` and the + surplus lanes are masked on every load and store, and the ``BLOCK_E`` segment + tail is masked by ``pos < end``. + +The two embedding GEMMs stay on cuBLAS (fp32 has no ``tl.dot`` tensor-core path) +and the ``ng x ng x 4`` Gram contraction that forms the final descriptor is +likewise left on cuBLAS; only the memory-bound tail is fused. + +The operator is a ``triton_op`` and composes under ``make_fx`` / ``torch.export`` +-- the graph lower is the pt_expt ``.pt2`` export target, whose forward traces +``forward_common_lower_graph`` over the edge schema. +""" + +from __future__ import ( + annotations, +) + +import logging + +import torch +import torch.nn.functional as F +from torch import ( + Tensor, +) +from torch.library import ( + triton_op, + wrap_triton, +) + +from deepmd.kernels.autotune import ( + register_autotuner, +) +from deepmd.kernels.triton.dpa1.activation import ( + TRITON_AVAILABLE, +) +from deepmd.kernels.triton.dpa1.tile_configs import ( + has_edge_config, + register_edge_config, + resolve_edge_config, +) +from deepmd.kernels.utils import ( + triton_infer_level, +) + +log = logging.getLogger(__name__) + +__all__ = [ + "concat_gate_placeholders", + "edge_conv", +] + + +# ====================================================================== +# Eager reference / fallback implementation +# ====================================================================== +def _edge_conv_reference( + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + dst: Tensor, + edge_mask: Tensor, + n_node: int, + resnet_mult: int, + act: int, + gated: int, +) -> Tensor: + """Eager ground truth for the graph environment convolution. + + ``gated`` selects the tebd-input mode: ``1`` (strip) applies the type-pair + gate ``gg = h2 * (1 + tt[idx] * sw)``; ``0`` (concat) skips it (the type + feature enters upstream through the embedding input), so ``tt`` / ``idx`` / + ``sw`` are ignored. Returns the unnormalized moment ``gr`` (n_node, 4, ng). + """ + activated = torch.tanh(z2) if act == 0 else F.silu(z2) + h2 = activated * idt + if resnet_mult == 2: + h2 = h2 + torch.cat([h1, h1], dim=-1) + elif resnet_mult == 1: + h2 = h2 + h1 + if gated: + gg_t = tt.index_select(0, idx.reshape(-1)) # (E, ng) + h2 = h2 * (1.0 + gg_t * sw[:, None]) + gg = h2 * edge_mask.to(z2.dtype)[:, None] # (E, ng) + # outer product (E, 4, ng): rr[:, k] * gg[:, c] + outer = rr[:, :, None] * gg[:, None, :] + gr = torch.zeros((n_node, 4, z2.shape[-1]), dtype=z2.dtype, device=z2.device) + gr.index_add_(0, dst, outer) + return gr + + +def _edge_conv_reference_backward( + grad_gr: Tensor, + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + dst: Tensor, + edge_mask: Tensor, + n_node: int, + resnet_mult: int, + act: int, + gated: int, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Closed-form gradient of :func:`_edge_conv_reference` w.r.t. the + coordinate-bearing inputs ``(z2, h1, sw, rr)``. + + A closed form (rather than a nested ``torch.autograd.grad``) is used so the + fallback composes under ``make_fx`` / ``torch.export``; the expressions + mirror the Triton backward kernel. ``gated == 0`` (concat) drops the type + gate, so ``fac == 1`` and ``grad_sw == 0``. + """ + if act == 0: + a = torch.tanh(z2) + act_grad = 1.0 - a * a + else: + s = torch.sigmoid(z2) + a = z2 * s + act_grad = s * (1.0 + z2 * (1.0 - s)) + h2 = a * idt + if resnet_mult == 2: + h2 = h2 + torch.cat([h1, h1], dim=-1) + elif resnet_mult == 1: + h2 = h2 + h1 + m = edge_mask.to(z2.dtype)[:, None] + if gated: + gg_t = tt.index_select(0, idx.reshape(-1)) # (E, ng) + fac = 1.0 + gg_t * sw[:, None] # d gg / d h2 + else: + fac = torch.ones_like(h2) + gg = h2 * fac * m + # gather the upstream moment gradient at each edge's center: (E, 4, ng) + g_e = grad_gr.index_select(0, dst) + # d gg = sum_k rr[k] * grad_gr[dst, k]; grad_rr[k] = sum_c gg[c] * grad_gr[dst, k] + dgg = (g_e * rr[:, :, None]).sum(dim=1) # (E, ng) + grad_rr = (g_e * gg[:, None, :]).sum(dim=-1) # (E, 4) + grad_h2 = dgg * fac * m + grad_z2 = grad_h2 * idt * act_grad + # fac = 1 + gg_t * sw, so d gg / d sw = h2 * gg_t (zero without the gate). + if gated: + grad_sw = (dgg * h2 * gg_t).sum(dim=-1) * m[:, 0] + else: + grad_sw = torch.zeros_like(sw) + if resnet_mult == 2: + h1d = h1.shape[-1] + grad_h1 = grad_h2[:, :h1d] + grad_h2[:, h1d : 2 * h1d] + elif resnet_mult == 1: + grad_h1 = grad_h2 + else: + grad_h1 = torch.zeros_like(h1) + return grad_z2, grad_h1, grad_sw, grad_rr + + +# ====================================================================== +# Triton kernels +# ====================================================================== +if TRITON_AVAILABLE: + import triton + import triton.language as tl + + from deepmd.kernels.triton.dpa1.activation import ( + activation, + activation_grad, + ) + + @triton.jit + def _edge_conv_fwd_kernel( + z2_ptr, # (E, NG) last embedding pre-activation + h1_ptr, # (E, H1) penultimate activation (resnet source) + idt_ptr, # (NG,) last-layer timestep + tt_ptr, # (P, NG) type-pair embedding table (strip gate) + idx_ptr, # (E,) per-edge type-pair row index (strip gate) + sw_ptr, # (E,) per-edge smooth cutoff (strip gate) + rr_ptr, # (E, 4) per-edge environment matrix + mask_ptr, # (E,) valid-edge mask (float 1/0) + order_ptr, # (E,) edge permutation grouping edges by center (dst-sorted) + segptr_ptr, # (N + 1,) CSR offsets of each node's segment in ``order`` + gr_ptr, # (N, 4, NG) accumulated moments + NG: tl.constexpr, + NGP: tl.constexpr, + H1: tl.constexpr, + RESNET_MULT: tl.constexpr, + ACT: tl.constexpr, + GATED: tl.constexpr, + BE: tl.constexpr, + ): + """Reduce one node's four moment rows over its CSR edge segment. + + One program per node: it streams the node's edges -- contiguous in the + dst-sorted ``order`` between ``segptr[node]`` and ``segptr[node + 1]`` -- + in ``BE``-edge blocks, forms ``gg`` in registers and accumulates the four + moment rows, writing ``gr[node]`` exactly once. This is the + contention-free analogue of an atomic scatter (whose ~nnei-way ``dst`` + collisions serialize): the reduction is register-local, like the dense + ``se_conv``, but streams only real edges (no ``sel`` padding). Edge + indices from ``order`` are int64 (``e * NG`` overflows int32 on large + systems); the channel mask ``rc < NG`` guards the padded width and the + segment mask ``pos < end`` guards the ``BE``-block tail. + + ``GATED`` selects the tebd-input mode: ``1`` (strip) applies the + type-pair gate ``gg = h2 * (1 + tt[idx] * sw)`` (the type feature + gathered inline through ``idx``); ``0`` (concat) skips it (``gg = h2``), + since concat feeds the type feature through the embedding input. + """ + node = tl.program_id(0) + rc = tl.arange(0, NGP) + cm = rc < NG + idt = tl.load(idt_ptr + rc, mask=cm, other=0.0) + beg = tl.load(segptr_ptr + node) + end = tl.load(segptr_ptr + node + 1) + acc_ty = z2_ptr.dtype.element_ty + acc0 = tl.zeros((NGP,), dtype=acc_ty) + acc1 = tl.zeros((NGP,), dtype=acc_ty) + acc2 = tl.zeros((NGP,), dtype=acc_ty) + acc3 = tl.zeros((NGP,), dtype=acc_ty) + for i0 in range(beg, end, BE): + pos = i0 + tl.arange(0, BE) + seg = pos < end + e = tl.load(order_ptr + pos, mask=seg, other=0).to(tl.int64) + m = seg[:, None] & cm[None, :] + z2 = tl.load(z2_ptr + e[:, None] * NG + rc[None, :], mask=m, other=0.0) + emask = tl.load(mask_ptr + e, mask=seg, other=0.0) + h2 = activation(z2, ACT) * idt[None, :] + if RESNET_MULT > 0: + # concat[h1, h1] (doubling) and identity share ``c mod H1``. + h1d = tl.load( + h1_ptr + e[:, None] * H1 + (rc % H1)[None, :], mask=m, other=0.0 + ) + h2 = h2 + h1d + if GATED: + idxe = tl.load(idx_ptr + e, mask=seg, other=0) + ggt = tl.load( + tt_ptr + idxe[:, None] * NG + rc[None, :], mask=m, other=0.0 + ) + swe = tl.load(sw_ptr + e, mask=seg, other=0.0) + h2 = h2 * (1.0 + ggt * swe[:, None]) + gg = h2 * emask[:, None] + base = e * 4 + s = tl.load(rr_ptr + base + 0, mask=seg, other=0.0) + rx = tl.load(rr_ptr + base + 1, mask=seg, other=0.0) + ry = tl.load(rr_ptr + base + 2, mask=seg, other=0.0) + rz = tl.load(rr_ptr + base + 3, mask=seg, other=0.0) + acc0 += tl.sum(s[:, None] * gg, axis=0) + acc1 += tl.sum(rx[:, None] * gg, axis=0) + acc2 += tl.sum(ry[:, None] * gg, axis=0) + acc3 += tl.sum(rz[:, None] * gg, axis=0) + ob = node * (4 * NG) + tl.store(gr_ptr + ob + 0 * NG + rc, acc0, mask=cm) + tl.store(gr_ptr + ob + 1 * NG + rc, acc1, mask=cm) + tl.store(gr_ptr + ob + 2 * NG + rc, acc2, mask=cm) + tl.store(gr_ptr + ob + 3 * NG + rc, acc3, mask=cm) + + @triton.jit + def _edge_conv_bwd_kernel( + ggr_ptr, # (N, 4, NG) upstream gradient of the moments + z2_ptr, + h1_ptr, + idt_ptr, + tt_ptr, # (P, NG) type-pair embedding table (strip gate) + idx_ptr, # (E,) per-edge type-pair row index (strip gate) + sw_ptr, # (E,) per-edge smooth cutoff (strip gate) + rr_ptr, + dst_ptr, + mask_ptr, + dz2_ptr, # (E, NG) + dh1_ptr, # (E, H1); written only when RESNET_MULT > 0 + dsw_ptr, # (E,); the switch gradient, nonzero only when GATED + drr_ptr, # (E, 4) + n_edge, + NG: tl.constexpr, + NGP: tl.constexpr, + H1: tl.constexpr, + H1P: tl.constexpr, + RESNET_MULT: tl.constexpr, + ACT: tl.constexpr, + GATED: tl.constexpr, + BE: tl.constexpr, + ): + """Backward of the edge scatter for a block of ``BE`` edges. + + The upstream moment gradient is gathered from ``grad_gr[dst]`` (a per-row + indirect load, no atomics); ``gg`` is recomputed in registers. The + residual gradient folds the ``ng // H1`` residual copies -- by a reshape + reduction when ``H1`` is a power of two, or by recomputing the two halves + at columns ``rj`` and ``rj + H1`` when ``H1`` is padded (the padded + reshape would split at ``H1P`` rather than the true ``H1``). + + When ``GATED`` (strip), ``gg = h2 * fac`` with ``fac = 1 + tt[idx] * sw``: + the gate multiplies ``grad_h2`` (hence ``dz2`` and the residual fold) and + contributes the switch gradient ``d L / d sw = sum_c dgg * h2 * tt[idx]``. + The padded-doubling fold gathers the gate for the two halves separately. + """ + pid = tl.program_id(0) + e = (pid * BE + tl.arange(0, BE)).to(tl.int64) + valid = e < n_edge + rc = tl.arange(0, NGP) + cm = rc < NG + m = valid[:, None] & cm[None, :] + idt = tl.load(idt_ptr + rc, mask=cm, other=0.0) + z2 = tl.load(z2_ptr + e[:, None] * NG + rc[None, :], mask=m, other=0.0) + emask = tl.load(mask_ptr + e, mask=valid, other=0.0) + a, ad = activation_grad(z2, ACT) + h2 = a * idt[None, :] + if RESNET_MULT > 0: + h1d = tl.load( + h1_ptr + e[:, None] * H1 + (rc % H1)[None, :], mask=m, other=0.0 + ) + h2 = h2 + h1d + if GATED: + idxe = tl.load(idx_ptr + e, mask=valid, other=0) + ggt = tl.load(tt_ptr + idxe[:, None] * NG + rc[None, :], mask=m, other=0.0) + swe = tl.load(sw_ptr + e, mask=valid, other=0.0) + fac = 1.0 + ggt * swe[:, None] + else: + fac = 1.0 + 0.0 * h2 + gg = h2 * fac * emask[:, None] + d = tl.load(dst_ptr + e, mask=valid, other=0) + ob = d[:, None] * (4 * NG) + rc[None, :] + g0 = tl.load(ggr_ptr + ob + 0 * NG, mask=m, other=0.0) + g1 = tl.load(ggr_ptr + ob + 1 * NG, mask=m, other=0.0) + g2 = tl.load(ggr_ptr + ob + 2 * NG, mask=m, other=0.0) + g3 = tl.load(ggr_ptr + ob + 3 * NG, mask=m, other=0.0) + base = e * 4 + s = tl.load(rr_ptr + base + 0, mask=valid, other=0.0) + rx = tl.load(rr_ptr + base + 1, mask=valid, other=0.0) + ry = tl.load(rr_ptr + base + 2, mask=valid, other=0.0) + rz = tl.load(rr_ptr + base + 3, mask=valid, other=0.0) + # d gg = sum_k rr[k] * grad_gr[dst, k] + dgg = s[:, None] * g0 + rx[:, None] * g1 + ry[:, None] * g2 + rz[:, None] * g3 + grad_h2 = dgg * fac * emask[:, None] + tl.store( + dz2_ptr + e[:, None] * NG + rc[None, :], grad_h2 * idt[None, :] * ad, mask=m + ) + # grad_rr[k] = sum_c gg[c] * grad_gr[dst, k, c] + tl.store(drr_ptr + base + 0, tl.sum(gg * g0, axis=1), mask=valid) + tl.store(drr_ptr + base + 1, tl.sum(gg * g1, axis=1), mask=valid) + tl.store(drr_ptr + base + 2, tl.sum(gg * g2, axis=1), mask=valid) + tl.store(drr_ptr + base + 3, tl.sum(gg * g3, axis=1), mask=valid) + if GATED: + # d L / d sw = sum_c dgg * (d gg / d sw) = sum_c dgg * h2 * tt[idx]. + tl.store( + dsw_ptr + e, + tl.sum(dgg * h2 * ggt, axis=1) * emask, + mask=valid, + ) + if RESNET_MULT == 1: + # Identity residual: grad_h1 equals grad_h2 (ng == H1). + tl.store(dh1_ptr + e[:, None] * H1 + rc[None, :], grad_h2, mask=m) + elif RESNET_MULT == 2 and H1P == H1: + grad_h1 = tl.sum(tl.reshape(grad_h2, (BE, 2, H1)), axis=1) + hj = tl.arange(0, H1) + tl.store( + dh1_ptr + e[:, None] * H1 + hj[None, :], + grad_h1, + mask=valid[:, None], + ) + elif RESNET_MULT == 2: + # Padded doubling: recompute the two halves at rj and rj + H1. + rj = tl.arange(0, H1P) + hm = valid[:, None] & (rj < H1)[None, :] + olo = d[:, None] * (4 * NG) + rj[None, :] + ohi = olo + H1 + g0lo = tl.load(ggr_ptr + olo + 0 * NG, mask=hm, other=0.0) + g1lo = tl.load(ggr_ptr + olo + 1 * NG, mask=hm, other=0.0) + g2lo = tl.load(ggr_ptr + olo + 2 * NG, mask=hm, other=0.0) + g3lo = tl.load(ggr_ptr + olo + 3 * NG, mask=hm, other=0.0) + g0hi = tl.load(ggr_ptr + ohi + 0 * NG, mask=hm, other=0.0) + g1hi = tl.load(ggr_ptr + ohi + 1 * NG, mask=hm, other=0.0) + g2hi = tl.load(ggr_ptr + ohi + 2 * NG, mask=hm, other=0.0) + g3hi = tl.load(ggr_ptr + ohi + 3 * NG, mask=hm, other=0.0) + dgg_lo = ( + s[:, None] * g0lo + + rx[:, None] * g1lo + + ry[:, None] * g2lo + + rz[:, None] * g3lo + ) + dgg_hi = ( + s[:, None] * g0hi + + rx[:, None] * g1hi + + ry[:, None] * g2hi + + rz[:, None] * g3hi + ) + if GATED: + # The gate factor differs per half; gather both columns. + ggt_lo = tl.load( + tt_ptr + idxe[:, None] * NG + rj[None, :], mask=hm, other=0.0 + ) + ggt_hi = tl.load( + tt_ptr + idxe[:, None] * NG + (rj + H1)[None, :], mask=hm, other=0.0 + ) + grad_h1 = ( + dgg_lo * (1.0 + ggt_lo * swe[:, None]) + + dgg_hi * (1.0 + ggt_hi * swe[:, None]) + ) * emask[:, None] + else: + grad_h1 = (dgg_lo + dgg_hi) * emask[:, None] + tl.store(dh1_ptr + e[:, None] * H1 + rj[None, :], grad_h1, mask=hm) + + +# ====================================================================== +# Dispatch, operator registration and public API +# ====================================================================== +def _use_triton(tensor: Tensor) -> bool: + return ( + TRITON_AVAILABLE + and tensor.is_cuda + and tensor.dtype in (torch.float32, torch.float64) + ) + + +def _edge_conv_fwd_impl( + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + dst: Tensor, + edge_mask: Tensor, + n_node: int, + resnet_mult: int, + act: int, + gated: int, + block_e: int, + num_warps: int, +) -> Tensor: + if not _use_triton(z2): + return _edge_conv_reference( + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dst, + edge_mask, + n_node, + resnet_mult, + act, + gated, + ) + _, ng = z2.shape + ngp = triton.next_power_of_2(ng) + # Destination-sorted CSR topology, built inside the op (integer, gradient- + # free): ``order`` groups edges by center, ``segptr`` gives each node's + # segment offsets. The node-parallel reduction is contention-free -- no + # atomics -- unlike the per-edge scatter it replaces. + order = torch.argsort(dst) + boundaries = torch.arange(n_node + 1, device=dst.device, dtype=dst.dtype) + segptr = torch.searchsorted(dst.index_select(0, order), boundaries).to(torch.int64) + gr = torch.empty((n_node, 4, ng), dtype=z2.dtype, device=z2.device) + wrap_triton(_edge_conv_fwd_kernel)[(n_node,)]( + z2, + h1, + idt, + tt, + idx, + sw, + rr, + edge_mask.to(z2.dtype), + order, + segptr, + gr, + NG=ng, + NGP=ngp, + H1=h1.shape[-1], + RESNET_MULT=resnet_mult, + ACT=act, + GATED=gated, + BE=block_e, + num_warps=num_warps, + ) + return gr + + +def _edge_conv_bwd_impl( + grad_gr: Tensor, + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + dst: Tensor, + edge_mask: Tensor, + n_node: int, + resnet_mult: int, + act: int, + gated: int, + block_e: int, + num_warps: int, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + if not _use_triton(z2): + return _edge_conv_reference_backward( + grad_gr, + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dst, + edge_mask, + n_node, + resnet_mult, + act, + gated, + ) + n_edge, ng = z2.shape + ngp = triton.next_power_of_2(ng) + dz2 = torch.empty_like(z2) + dh1 = torch.empty_like(h1) if resnet_mult > 0 else torch.zeros_like(h1) + dsw = torch.empty_like(sw) if gated else torch.zeros_like(sw) + drr = torch.empty_like(rr) + grid = (triton.cdiv(n_edge, block_e),) + wrap_triton(_edge_conv_bwd_kernel)[grid]( + grad_gr.contiguous(), + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dst, + edge_mask.to(z2.dtype), + dz2, + dh1, + dsw, + drr, + n_edge, + NG=ng, + NGP=ngp, + H1=h1.shape[-1], + H1P=triton.next_power_of_2(h1.shape[-1]), + RESNET_MULT=resnet_mult, + ACT=act, + GATED=gated, + BE=block_e, + num_warps=num_warps, + ) + return dz2, dh1, dsw, drr + + +_edge_conv_op = triton_op("dpa1_triton::edge_conv", mutates_args=())( + _edge_conv_fwd_impl +) +_edge_conv_bwd_op = triton_op("dpa1_triton::edge_conv_bwd", mutates_args=())( + _edge_conv_bwd_impl +) + + +@_edge_conv_op.register_fake +def _( + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dst, + edge_mask, + n_node, + resnet_mult, + act, + gated, + block_e, + num_warps, +): + return z2.new_empty((n_node, 4, z2.shape[1])) + + +@_edge_conv_bwd_op.register_fake +def _( + grad_gr, + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dst, + edge_mask, + n_node, + resnet_mult, + act, + gated, + block_e, + num_warps, +): + return ( + torch.empty_like(z2), + torch.empty_like(h1), + torch.empty_like(sw), + torch.empty_like(rr), + ) + + +def _edge_conv_setup_context(ctx, inputs, output): + ( + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dst, + edge_mask, + n_node, + resnet_mult, + act, + gated, + block_e, + num_warps, + ) = inputs + ctx.save_for_backward(z2, h1, idt, tt, idx, sw, rr, dst, edge_mask) + ctx.n_node = n_node + ctx.resnet_mult = resnet_mult + ctx.act = act + ctx.gated = gated + ctx.block_e = block_e + ctx.num_warps = num_warps + + +def _edge_conv_backward(ctx, grad_gr): + z2, h1, idt, tt, idx, sw, rr, dst, edge_mask = ctx.saved_tensors + grad_z2, grad_h1, grad_sw, grad_rr = _edge_conv_bwd_op( + grad_gr.contiguous(), + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dst, + edge_mask, + ctx.n_node, + ctx.resnet_mult, + ctx.act, + ctx.gated, + ctx.block_e, + ctx.num_warps, + ) + # ``idt`` / ``tt`` / ``idx`` / ``dst`` / ``edge_mask`` do not depend on + # coordinates; ``n_node``, ``resnet_mult``, ``act``, ``gated`` and the launch + # config are static. + return ( + grad_z2, + grad_h1, + None, + None, + None, + grad_sw, + grad_rr, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +_edge_conv_op.register_autograd( + _edge_conv_backward, setup_context=_edge_conv_setup_context +) + + +def concat_gate_placeholders(z2: Tensor, ng: int) -> tuple[Tensor, Tensor, Tensor]: + """Unused ``(tt, idx, sw)`` gate inputs for the concat (``gated == 0``) call. + + Concat carries the type feature through the embedding input, so the strip + type-pair gate is inactive and :func:`edge_conv` skips these tensors. They + still fill the operator signature; minimal (single-element) placeholders + suffice and keep the traced graph lean. + + Parameters + ---------- + z2 : Tensor + A tensor sharing the target device and floating dtype. + ng : int + Embedding channel width (the ``tt`` table's column count). + + Returns + ------- + tuple[Tensor, Tensor, Tensor] + The placeholder ``(tt, idx, sw)``: ``tt`` (1, ng), ``idx`` (1,) int64, + ``sw`` (1,), matching the per-edge switch layout. + """ + tt = z2.new_zeros(1, ng) + idx = torch.zeros(1, dtype=torch.long, device=z2.device) + sw = z2.new_ones(1) + return tt, idx, sw + + +def edge_conv( + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + dst: Tensor, + edge_mask: Tensor, + n_node: int, + resnet_mult: int, + act: int, + gated: int, +) -> Tensor: + """Fused graph-native ``se_atten`` environment convolution (attn-free). + + Parameters + ---------- + z2 : Tensor + Last embedding pre-activation with shape (E, ng), where ``E`` is the + edge count. + h1 : Tensor + Penultimate embedding activation with shape (E, H1); read only when + ``resnet_mult > 0``, where ``ng == resnet_mult * H1``. + idt : Tensor + Last-layer per-channel timestep with shape (ng,); ones without + ``resnet_dt``. + tt : Tensor + Type-pair embedding table with shape (P, ng); gathered per edge by + ``idx`` for the strip gate. A placeholder when ``gated == 0``. + idx : Tensor + Per-edge row index into ``tt`` with shape (E,) (strip). One-side uses the + neighbor type; two-side folds the ``(center, neighbor)`` pair. A + placeholder when ``gated == 0``. + sw : Tensor + Per-edge smooth cutoff with shape (E,) (strip); the switch multiplying + the gate, or ones when the descriptor is non-smooth. A placeholder when + ``gated == 0``. + rr : Tensor + Per-edge environment matrix ``(s, s*x/r, s*y/r, s*z/r)`` with shape + (E, 4). + dst : Tensor + Center-node index of each edge with shape (E,); the ``segment_sum`` + segment id (arbitrary order). + edge_mask : Tensor + Valid-edge mask with shape (E,); padding edges are dropped from the + reduction (``edge_env_mat`` leaves them nonzero). + n_node : int + Number of nodes ``N`` (the segment count). + resnet_mult : int + Residual structure of the last layer: ``2`` (width doubling), ``1`` + (identity), or ``0`` (no residual). + act : int + Last-layer activation code from :data:`.activation.ACT_CODES`: ``0`` for + ``tanh``, ``1`` for ``silu``. + gated : int + Tebd-input mode: ``1`` applies the strip type-pair gate + ``gg = h2 * (1 + tt[idx] * sw)``; ``0`` (concat) leaves ``gg = h2`` (the + type feature entered the embedding input upstream). + + Returns + ------- + Tensor + Unnormalized moments ``gr`` with shape (N, 4, ng), equal to the + ``segment_sum`` of ``rr[:, k] * gg[:, c]`` over edge centers. The + ``1 / nnei`` normalization is applied by the descriptor afterwards. + + Notes + ----- + The launch configuration ``(BLOCK_E, num_warps)`` -- the forward's per-node + segment-block width and the warp count -- is resolved from the active + ``DP_TRITON_INFER`` level via :func:`.tile_configs.resolve_edge_config`. The + operator composes under ``make_fx`` / ``torch.compile`` and exposes an + inference force gradient through the registered backward, so it is baked into + the pt_expt graph-form ``.pt2`` export. + """ + block_e, num_warps = resolve_edge_config( + int(z2.shape[-1]), int(h1.shape[-1]), triton_infer_level() + ) + return _edge_conv_op( + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dst, + edge_mask, + n_node, + resnet_mult, + act, + gated, + block_e, + num_warps, + ) + + +def _autotune_edge(model: torch.nn.Module, level: int, device: torch.device) -> None: + """Sweep the ``edge_conv`` launch table for a model about to be frozen. + + Collects the ``(ng, H1)`` shape key of every attention-free ``se_atten`` + descriptor in ``model`` (the graph lower's ``edge_conv`` users, in either + concat or strip tebd-input mode) and sweeps the keys the built-in / + freeze-time tables do not yet cover on the target GPU, registering the + winners so the ``resolve_edge_config`` lookups made while tracing bake tuned + launches into the graph-form ``.pt2``. Keys already covered cost nothing. + """ + from deepmd.kernels.triton.dpa1.sweep_tile_configs import ( + sweep_edge, + ) + + device_name = torch.cuda.get_device_name(device) + keys: set[tuple[int, int]] = set() + for module in model.modules(): + eligible = getattr(module, "_fused_eligible", None) + se = getattr(module, "se_atten", None) + if not (callable(eligible) and eligible("triton")): + continue + if se is None or se.tebd_input_mode not in ("concat", "strip"): + continue + weight = se.embeddings[0].layers[-1].w + keys.add((int(weight.shape[1]), int(weight.shape[0]))) + tuned: dict[tuple[int, int], tuple[int, int]] = {} + for ng, h1 in sorted(keys): + if has_edge_config(ng, h1): + continue + config = sweep_edge(ng, h1, device=device) + register_edge_config(device_name, ng, h1, config) + tuned[(ng, h1)] = config + if tuned: + log.info("DPA1 edge_conv: tuned launch configs %s on %s.", tuned, device_name) + else: + log.info( + "DPA1 edge_conv: launch table already covers this checkpoint's shapes " + "on %s; no tuning needed.", + device_name, + ) + + +if TRITON_AVAILABLE: + register_autotuner("dpa1_edge_conv", _autotune_edge) diff --git a/deepmd/kernels/triton/dpa1/gemm_fp16x3.py b/deepmd/kernels/triton/dpa1/gemm_fp16x3.py new file mode 100644 index 0000000000..15e5ff5c2c --- /dev/null +++ b/deepmd/kernels/triton/dpa1/gemm_fp16x3.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +# pyright: reportMissingImports=false +# ruff: noqa: ANN001, ANN202 +"""fp16x3 split-compensated dense GEMM for the DPA1 embedding (``DP_TRITON_INFER +>= 3``). + +The DPA1 embedding GEMMs run on a large edge/row count ``M`` (hundreds of +thousands) with modest ``K`` / ``N``; the profile shows they are compute-bound +(arithmetic intensity well above the H20 fp32 compute/bandwidth break-even), so +the tensor cores can beat the fp32 FFMA path that cuBLAS uses. This operator +evaluates ``C = A @ B`` as three fp16 tensor-core products with fp32 +accumulation (a two-term Ootomo split):: + + A = A_hi + A_lo, B = B_hi + B_lo (fp16 head + fp16 tail) + C ~= A_hi B_hi + A_hi B_lo + A_lo B_hi (the A_lo B_lo term, ~2^-22 + relative, is dropped) + +An fp16 multiply into an fp32 accumulator is exact, so the only error is the two +split truncations. The head product and the two tail corrections accumulate in +*separate* fp32 accumulators merged once per tile -- chaining all three into one +absorbs the small tails against the large head partial sums and doubles the +error. Tails are stored pre-scaled by ``2^11`` (the fp16 mantissa width) so an +element below the fp16 subnormal range keeps its correction; the epilogue scales +back. Measured on the DPA1 embedding shapes the maximum relative error matches +the fp32 reference (~3e-7) at ~1.6x the cuBLAS fp32 throughput. + +Only the compute-bound, large-``M`` embedding GEMMs are eligible; the small +fitting GEMMs (``M`` ~ nloc) stay on cuBLAS, where fp16x3 loses to the launch and +split overhead. The operator is inference-only (level 3) and returns the input +gradient (the coordinate force path); the weight gradient is not formed +(training keeps the fp32 reference path). + +The launch pins ``num_stages = 1``. A split-fp16 k-loop is a known failure mode +of the Triton software pipeliner: at ``num_stages >= 2`` the prefetch reorders +the head/tail split across loop iterations and can emit silent ``NaN`` on some +shapes, and -- as the SeZM stacked-GEMM tuning established -- no finite sample +certifies every problem size, so a pipelined config validated at one row count +may still poison another. Disabling the pipeliner removes the reordering +entirely and is therefore structurally finite for all shapes; on the DPA1 +embedding widths it costs under 20% of the GEMM (still ~1.3x over cuBLAS fp32), +which is negligible end to end. A faster stage count would require the SeZM-style +per-shape finiteness sweep and is not warranted for this single GEMM. +""" + +from __future__ import ( + annotations, +) + +import torch +from torch import ( + Tensor, +) +from torch.library import ( + triton_op, + wrap_triton, +) + +from deepmd.kernels.triton.dpa1.activation import ( + TRITON_AVAILABLE, +) +from deepmd.kernels.utils import ( + triton_infer_level, +) + +__all__ = [ + "embed_gemm_fp16x3", + "embed_last_gemm", +] + +# fp16 mantissa-width tail prescale (2^11); an exact power of two, so the split +# and its inverse are lossless. +_TAIL_SCALE = 2048.0 + +if TRITON_AVAILABLE: + import triton + import triton.language as tl + + @triton.jit + def _gemm_fp16x3_kernel( + a_ptr, # (M, K) row-major + b_ptr, # (K, N) row-major + c_ptr, # (M, N) row-major + M, + N, + K, + S: tl.constexpr, # tail prescale + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + ): + """``C = A @ B`` via the two-term Ootomo split on fp16 tensor cores. + + The head (``A_hi B_hi``) and the two tail corrections (``A_hi B_lo`` and + ``A_lo B_hi``) accumulate in separate fp32 tiles, merged once after the + k-loop; the dropped ``A_lo B_lo`` term is ~``2^-22`` relative. + """ + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + # Row offsets are int64: M is the edge count, and the flat element + # offsets rm * K / rm * N exceed int32 beyond ~17M rows at N = 128. + rm = (pid_m * BM + tl.arange(0, BM)).to(tl.int64) + rn = pid_n * BN + tl.arange(0, BN) + rk = tl.arange(0, BK) + a_ptrs = a_ptr + rm[:, None] * K + rk[None, :] + b_ptrs = b_ptr + rk[:, None] * N + rn[None, :] + m_row = rm[:, None] < M + n_col = rn[None, :] < N + acc_hi = tl.zeros((BM, BN), dtype=tl.float32) + acc_lo = tl.zeros((BM, BN), dtype=tl.float32) + for k0 in range(0, K, BK): + kk = rk + k0 + a = tl.load(a_ptrs, mask=m_row & (kk[None, :] < K), other=0.0) + b = tl.load(b_ptrs, mask=(kk[:, None] < K) & n_col, other=0.0) + a_hi = a.to(tl.float16) + a_lo = ((a - a_hi.to(tl.float32)) * S).to(tl.float16) + b_hi = b.to(tl.float16) + b_lo = ((b - b_hi.to(tl.float32)) * S).to(tl.float16) + acc_hi += tl.dot(a_hi, b_hi, out_dtype=tl.float32, input_precision="ieee") + acc_lo += tl.dot(a_hi, b_lo, out_dtype=tl.float32, input_precision="ieee") + acc_lo += tl.dot(a_lo, b_hi, out_dtype=tl.float32, input_precision="ieee") + a_ptrs += BK + b_ptrs += BK * N + acc = acc_hi + acc_lo / S + c_ptrs = c_ptr + rm[:, None] * N + rn[None, :] + tl.store(c_ptrs, acc, mask=m_row & n_col) + + +def _use_triton(tensor: Tensor) -> bool: + return TRITON_AVAILABLE and tensor.is_cuda and tensor.dtype == torch.float32 + + +def _gemm_impl( + a: Tensor, b: Tensor, bm: int, bn: int, bk: int, num_warps: int, num_stages: int +) -> Tensor: + # a: (M, K) contiguous, b: (K, N) contiguous. + if not _use_triton(a): + return a @ b + m, k = a.shape + n = b.shape[1] + c = torch.empty((m, n), dtype=torch.float32, device=a.device) + grid = (triton.cdiv(m, bm), triton.cdiv(n, bn)) + wrap_triton(_gemm_fp16x3_kernel)[grid]( + a.contiguous(), + b.contiguous(), + c, + m, + n, + k, + S=_TAIL_SCALE, + BM=bm, + BN=bn, + BK=bk, + num_warps=num_warps, + num_stages=num_stages, + ) + return c + + +_gemm_op = triton_op("dpa1_triton::embed_gemm_fp16x3", mutates_args=())(_gemm_impl) + + +@_gemm_op.register_fake +def _(a, b, bm, bn, bk, num_warps, num_stages): + return a.new_empty((a.shape[0], b.shape[1])) + + +def _gemm_setup_context(ctx, inputs, output): + a, b, bm, bn, bk, num_warps, num_stages = inputs + ctx.save_for_backward(b) + ctx.cfg = (bm, bn, bk, num_warps, num_stages) + + +def _gemm_backward(ctx, grad_c): + (b,) = ctx.saved_tensors + bm, bn, bk, num_warps, num_stages = ctx.cfg + # d(C = A @ B) / dA = grad_C @ B^T; B^T is (N, K), materialized contiguous + # (B is tiny -- the embedding weight). The weight gradient is not formed + # (inference-only; training uses the fp32 reference path). + grad_a = _gemm_op( + grad_c.contiguous(), b.t().contiguous(), bm, bn, bk, num_warps, num_stages + ) + return grad_a, None, None, None, None, None, None + + +_gemm_op.register_autograd(_gemm_backward, setup_context=_gemm_setup_context) + + +def embed_gemm_fp16x3(a: Tensor, b: Tensor) -> Tensor: + """fp16x3 dense matmul ``a @ b`` for the compute-bound DPA1 embedding. + + Parameters + ---------- + a : Tensor + Left operand with shape (M, K); ``M`` is the (large) edge/row count. + b : Tensor + Right operand (embedding weight) with shape (K, N). + + Returns + ------- + Tensor + ``a @ b`` with shape (M, N), computed on fp16 tensor cores with the + two-term split (fp32-reference accuracy, ~2^-22 rounding). Falls back to + a plain fp32 matmul off the CUDA fp32 path. + + Notes + ----- + Inference-only (``DP_TRITON_INFER >= 3``); the registered backward returns + the input gradient (force path) and no weight gradient. Composes under + ``make_fx`` / ``torch.export`` as a ``triton_op``. + """ + bm, bn, bk, num_warps, num_stages = _resolve_gemm_config( + int(a.shape[1]), int(b.shape[1]) + ) + return _gemm_op(a, b, bm, bn, bk, num_warps, num_stages) + + +def embed_last_gemm(h: Tensor, w: Tensor, bias: Tensor | None) -> Tensor: + """Last embedding-layer pre-activation ``z2 = h @ w (+ bias)``. + + Routes through the fp16x3 tensor-core GEMM at ``DP_TRITON_INFER >= 3`` (the + large-``M`` embedding GEMM is compute-bound and beats cuBLAS fp32 there), + else a plain fp32 matmul. Shared by the strip (``se_atten_conv``) and concat + / graph (pt_expt) fused paths. + + Parameters + ---------- + h : Tensor + Penultimate embedding activation with shape (..., h1_dim): 2-D + ``(E, h1_dim)`` on the graph path, 3-D ``(nfnl, nnei, h1_dim)`` on the + dense path (the leading axes are flattened for the 2-D GEMM). + w : Tensor + Last-layer weight with shape (h1_dim, ng). + bias : Tensor or None + Optional last-layer bias with shape (ng,). + + Returns + ------- + Tensor + The pre-activation ``z2`` with shape (..., ng). + """ + if triton_infer_level() >= 3: + lead = h.shape[:-1] + flat = embed_gemm_fp16x3(h.reshape(-1, h.shape[-1]).contiguous(), w) + z2 = flat.reshape(*lead, w.shape[-1]) + else: + z2 = torch.matmul(h, w) + if bias is not None: + z2 = z2 + bias + return z2 + + +def _resolve_gemm_config(k: int, n: int) -> tuple[int, int, int, int, int]: + """Launch config ``(BM, BN, BK, num_warps, num_stages)`` for the embedding GEMM. + + Keyed by ``(K, N)`` (the small contracted/output widths that bound the tile); + ``M`` is streamed and does not shift the optimum. Unswept shapes take a + spill-safe default. ``num_stages`` is pinned to 1 across every entry: the + split-fp16 k-loop is miscompiled by the Triton pipeliner at higher stage + counts (silent ``NaN``), so only stage 1 is certifiably finite for all shapes. + """ + return _GEMM_CONFIGS.get((int(k), int(n)), _GEMM_DEFAULT) + + +# BM, BN, BK, num_warps, num_stages. Default is a safe 64x64x32 tile with the +# pipeliner disabled (num_stages=1) -- see the module docstring on fp16x3 NaN. +_GEMM_DEFAULT: tuple[int, int, int, int, int] = (64, 64, 32, 4, 1) +_GEMM_CONFIGS: dict[tuple[int, int], tuple[int, int, int, int, int]] = {} diff --git a/deepmd/kernels/triton/dpa1/se_conv.py b/deepmd/kernels/triton/dpa1/se_conv.py new file mode 100644 index 0000000000..96390baf86 --- /dev/null +++ b/deepmd/kernels/triton/dpa1/se_conv.py @@ -0,0 +1,849 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +# pyright: reportMissingImports=false +# ruff: noqa: ANN001, ANN202 +"""Fused environment convolution for the DPA1 (``se_atten``) descriptor. + +For the attention-free (``attn_layer == 0``), strip-embedding ``se_atten`` path +the per-edge descriptor tail is + + ``h2 = act(z2) * idt + resnet(h1)`` (last embedding layer + resnet) + ``gg = h2 * (1 + tt[idx] * sw)`` (type-pair gate + smooth cutoff) + ``xyz[k, c] = sum_j rr[j, k] * gg[j, c]`` (moment accumulation) + +where ``act`` is the last layer's activation (``tanh`` or ``silu``), ``z2`` is +the last embedding pre-activation ``h1 @ W2 + b2``, ``h1`` is the +penultimate activation feeding the resnet, ``idt`` is the last layer's per- +channel timestep (all ones when ``resnet_dt`` is off), ``tt`` is the type-pair +embedding table with per-edge row index ``idx``, ``sw`` is the smooth radial +cutoff, and ``rr`` is the ``(s, s*x/r, s*y/r, s*z/r)`` environment matrix. The +``1 / nnei`` normalization of the moment is applied by the descriptor after this +operator (matching the eager and tabulated paths), so the operator returns the +unnormalized moment ``rr^T @ gg``. + +The last-layer resnet takes one of three shapes, selected by ``resnet_mult`` +(``= ng // H1`` when the layer adds a residual, else ``0``): + +- ``resnet_mult == 2`` -- width doubling, ``resnet(h1)[c] = concat[h1, h1][c]``. +- ``resnet_mult == 1`` -- identity, ``resnet(h1) = h1`` (``ng == H1``). +- ``resnet_mult == 0`` -- no residual; ``h1`` does not enter the activation. + +The doubling and identity cases share one addressing rule ``h1[c mod H1]`` and +one backward fold (summing the ``ng // H1`` residual copies), so a single kernel +covers both; the no-residual case skips the ``h1`` read. + +The eager path materializes three ``(E, ng)`` tensors (``h2``, the gathered type +feature, and ``gg``) and then runs a batched ``rr^T @ gg`` matmul whose ``M = 4`` +contraction uses cuBLAS poorly. This operator fuses the whole tail into one +node-parallel kernel: each program owns one node, streams its neighbors, forms +``gg`` in registers (never materializing any ``(E, ng)`` tensor, and gathering +the type feature inline) and accumulates the four moment rows. The two embedding +GEMMs (``h0 @ W1`` and ``h1 @ W2``) stay on cuBLAS -- in the fp32 regime this +descriptor runs in, Triton ``tl.dot`` has no tensor-core path and cannot beat +cuBLAS, so only the memory-bound, non-GEMM tail is fused. The trailing +``ng x ng x 4`` Gram contraction that forms the final descriptor is likewise +left on cuBLAS. + +Design notes and pitfalls +------------------------- +- **Inlined activations (``tanh`` / ``silu``).** The last-layer activation is + inlined for both value and derivative, selected by the ``act`` code + (:data:`.activation.ACT_CODES`). Only these activations are eligible; any other + keeps the dense reference path (the descriptor gates on the activation name). + Both are smooth, so the inference force stays differentiable. +- **Input-precision accumulation.** The moment is accumulated in the input + dtype -- fp32 for the eager DPA1 path, fp64 for the float64 pt_expt / export + path -- and no ``tl.dot`` is used. +- **Arbitrary channel widths via padding.** ``tl.arange`` requires a + power-of-two bound, so the channel axis is carried at the padded width + ``next_pow2(ng)`` and the ``next_pow2(ng) - ng`` surplus lanes are masked to + zero on every load and skipped on every store. Padding to the next power of + two (rather than tiling the channel axis) keeps the kernel single-pass and + is uniformly faster than the eager path even at the worst padding ratio, + because the extra lanes are pure bandwidth with no cross-lane dependence. The + backward residual fold additionally depends on the true (unpadded) ``H1``; + see :func:`_se_conv_bwd_kernel`. +- **Backward register footprint.** The backward holds a ``(BLOCK_N, ng)`` block + live; an oversized ``BLOCK_N`` spills and collapses throughput far more + sharply than in the forward. Configurations come from + :func:`.tile_configs.resolve_conv_config`, whose default is deliberately + small. +- **Inference-only autograd.** The registered backward returns gradients for + the coordinate-bearing inputs (``z2``, ``h1``, ``sw``, ``rr``) that carry the + force; the timestep, type table and index do not depend on coordinates. + Training keeps the dense reference path (the gate is inference-only), so their + gradients are never required here. +""" + +from __future__ import ( + annotations, +) + +import logging + +import torch +import torch.nn.functional as F +from torch import ( + Tensor, +) +from torch.library import ( + triton_op, + wrap_triton, +) + +from deepmd.kernels.autotune import ( + register_autotuner, +) +from deepmd.kernels.triton.dpa1.activation import ( + ACT_CODES, + TRITON_AVAILABLE, +) +from deepmd.kernels.triton.dpa1.gemm_fp16x3 import ( + embed_last_gemm, +) +from deepmd.kernels.triton.dpa1.tile_configs import ( + has_conv_config, + register_conv_config, + resolve_conv_config, +) +from deepmd.kernels.utils import ( + triton_infer_level, +) + +log = logging.getLogger(__name__) + +__all__ = [ + "concat_gate_placeholders", + "se_atten_conv", + "se_conv", +] + + +# ====================================================================== +# Eager reference / fallback implementation +# ====================================================================== +def _se_conv_reference( + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + resnet_mult: int, + act: int, + gated: int, +) -> Tensor: + """Eager ground truth for the fused environment convolution. + + ``gated`` selects the tebd-input mode: ``1`` (strip) applies the type-pair + gate; ``0`` (concat) skips it (the type feature enters upstream through the + embedding input), so ``tt`` / ``idx`` / ``sw`` are ignored. + """ + nfnl, nnei, ng = z2.shape + activated = torch.tanh(z2) if act == 0 else F.silu(z2) + h2 = activated * idt + if resnet_mult == 2: + h2 = h2 + torch.cat([h1, h1], dim=-1) + elif resnet_mult == 1: + h2 = h2 + h1 + if gated: + gg_t = tt.index_select(0, idx.reshape(-1)).reshape(nfnl, nnei, ng) + gg = h2 * (1.0 + gg_t * sw.unsqueeze(-1)) + else: + gg = h2 + return torch.matmul(rr.transpose(1, 2), gg) + + +def _se_conv_reference_backward( + grad_xyz: Tensor, + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + resnet_mult: int, + act: int, + gated: int, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Closed-form gradient of :func:`_se_conv_reference` w.r.t. the coordinate- + bearing inputs ``(z2, h1, sw, rr)``. + + A closed form (rather than a nested ``torch.autograd.grad``) is used so the + fallback composes under ``make_fx`` / ``torch.export``: the tracer runs this + body with grad tracking disabled, where an inner autograd call would find no + graph. The expressions mirror the Triton backward kernel. ``gated == 0`` + (concat) drops the type gate, so ``fac == 1`` and ``grad_sw == 0``. + """ + nfnl, nnei, ng = z2.shape + if act == 0: + a = torch.tanh(z2) + act_grad = 1.0 - a * a + else: + s = torch.sigmoid(z2) + a = z2 * s + act_grad = s * (1.0 + z2 * (1.0 - s)) + h2 = a * idt + if resnet_mult == 2: + h2 = h2 + torch.cat([h1, h1], dim=-1) + elif resnet_mult == 1: + h2 = h2 + h1 + if gated: + gg_t = tt.index_select(0, idx.reshape(-1)).reshape(nfnl, nnei, ng) + fac = 1.0 + gg_t * sw.unsqueeze(-1) # d gg / d h2 + else: + fac = torch.ones_like(h2) + gg = h2 * fac + # d(xyz = rr^T @ gg): grad_gg = rr @ grad_xyz, grad_rr = gg @ grad_xyz^T. + grad_gg = torch.matmul(rr, grad_xyz) + grad_rr = torch.matmul(gg, grad_xyz.transpose(1, 2)) + grad_h2 = grad_gg * fac + grad_z2 = grad_h2 * idt * act_grad + # fac = 1 + gg_t * sw, so d gg / d sw = h2 * gg_t (zero without the gate). + if gated: + grad_sw = (grad_gg * h2 * gg_t).sum(dim=-1) + else: + grad_sw = torch.zeros_like(sw) + if resnet_mult == 2: + hh = h1.shape[-1] + grad_h1 = grad_h2[..., :hh] + grad_h2[..., hh : 2 * hh] + elif resnet_mult == 1: + grad_h1 = grad_h2 + else: + grad_h1 = torch.zeros_like(h1) + return grad_z2, grad_h1, grad_sw, grad_rr + + +# ====================================================================== +# Triton kernels +# ====================================================================== +if TRITON_AVAILABLE: + import triton + import triton.language as tl + + from deepmd.kernels.triton.dpa1.activation import ( + activation, + activation_grad, + ) + + @triton.jit + def _se_conv_fwd_kernel( + z2_ptr, # (N, NNEI, NG) last embedding pre-activation + h1_ptr, # (N, NNEI, H1) penultimate activation (resnet source) + idt_ptr, # (NG,) last-layer timestep (ones without resnet_dt) + tt_ptr, # (P, NG) type-pair embedding table + idx_ptr, # (N * NNEI,) per-edge type-pair row index + sw_ptr, # (N, NNEI) smooth radial cutoff + rr_ptr, # (N, NNEI, 4) environment matrix + out_ptr, # (N, 4, NG) accumulated moments + NNEI, + H1: tl.constexpr, + NG: tl.constexpr, + NGP: tl.constexpr, + RESNET_MULT: tl.constexpr, + ACT: tl.constexpr, + GATED: tl.constexpr, + BN: tl.constexpr, + ): + """Accumulate the four unnormalized moment rows over a node's neighbors. + + ``gg`` is formed in registers per neighbor block and never written to + global memory; the type feature is gathered inline through ``idx``. The + channel axis is carried at the padded width ``NGP = next_pow2(NG)`` (the + ``tl.arange`` bound must be a power of two); the ``NGP - NG`` padding + lanes are masked to zero on every load and skipped on the final store, + so they contribute nothing to the moment reduction. + + ``GATED`` selects the tebd-input mode: ``1`` (strip) applies the + type-pair gate ``gg = h2 * (1 + tt[idx] * sw)``; ``0`` (concat) skips it + (``gg = h2``), since concat feeds the type feature through the embedding + input instead of a multiplicative gate. + """ + node = tl.program_id(0) + rc = tl.arange(0, NGP) + # The channel mask is only material when the width is padded (NGP > NG); + # for a power-of-two width it is compile-time all-true. It is elided from + # the per-neighbor loads and the ``dz2`` store (the bandwidth-bound hot + # path) so an unpadded launch matches the plain neighbor-masked fast path + # exactly; the cheap once-per-node loads/stores keep it unconditionally. + cm = rc < NG + idt = tl.load(idt_ptr + rc, mask=cm, other=0.0) + # Accumulate in the input precision so the kernel serves both fp32 + # (the eager DPA1 path) and fp64 (the float64 pt_expt / export path). + acc_ty = z2_ptr.dtype.element_ty + acc0 = tl.zeros((NGP,), dtype=acc_ty) + acc1 = tl.zeros((NGP,), dtype=acc_ty) + acc2 = tl.zeros((NGP,), dtype=acc_ty) + acc3 = tl.zeros((NGP,), dtype=acc_ty) + for n0 in range(0, NNEI, BN): + offs = n0 + tl.arange(0, BN) + nmask = offs < NNEI + m = (nmask[:, None] & cm[None, :]) if NGP != NG else nmask[:, None] + e = node * NNEI + offs + z2 = tl.load(z2_ptr + e[:, None] * NG + rc[None, :], mask=m, other=0.0) + base = e * 4 + s = tl.load(rr_ptr + base + 0, mask=nmask, other=0.0) + rx = tl.load(rr_ptr + base + 1, mask=nmask, other=0.0) + ry = tl.load(rr_ptr + base + 2, mask=nmask, other=0.0) + rz = tl.load(rr_ptr + base + 3, mask=nmask, other=0.0) + h2 = activation(z2, ACT) * idt[None, :] + if RESNET_MULT > 0: + # concat[h1, h1] (doubling) and identity share one addressing + # rule; ``c mod H1`` reads the residual source for either width. + h1d = tl.load( + h1_ptr + e[:, None] * H1 + (rc % H1)[None, :], mask=m, other=0.0 + ) + h2 = h2 + h1d + if GATED: + idx = tl.load(idx_ptr + e, mask=nmask, other=0) + ggt = tl.load( + tt_ptr + idx[:, None] * NG + rc[None, :], mask=m, other=0.0 + ) + sw = tl.load(sw_ptr + e, mask=nmask, other=0.0) + gg = h2 * (1.0 + ggt * sw[:, None]) + else: + gg = h2 + acc0 += tl.sum(s[:, None] * gg, axis=0) + acc1 += tl.sum(rx[:, None] * gg, axis=0) + acc2 += tl.sum(ry[:, None] * gg, axis=0) + acc3 += tl.sum(rz[:, None] * gg, axis=0) + ob = node * 4 * NG + tl.store(out_ptr + ob + 0 * NG + rc, acc0, mask=cm) + tl.store(out_ptr + ob + 1 * NG + rc, acc1, mask=cm) + tl.store(out_ptr + ob + 2 * NG + rc, acc2, mask=cm) + tl.store(out_ptr + ob + 3 * NG + rc, acc3, mask=cm) + + @triton.jit + def _se_conv_bwd_kernel( + gout_ptr, # (N, 4, NG) upstream gradient of the moments + z2_ptr, + h1_ptr, + idt_ptr, + tt_ptr, + idx_ptr, + sw_ptr, + rr_ptr, + dz2_ptr, # (N, NNEI, NG) + dh1_ptr, # (N, NNEI, H1); written only when RESNET_MULT > 0 + dsw_ptr, # (N, NNEI) + drr_ptr, # (N, NNEI, 4) + NNEI, + H1: tl.constexpr, + NG: tl.constexpr, + NGP: tl.constexpr, + H1P: tl.constexpr, + RESNET_MULT: tl.constexpr, + ACT: tl.constexpr, + GATED: tl.constexpr, + BN: tl.constexpr, + ): + """Backward of the moment accumulation for one node. + + The upstream moment gradient is loaded once per node and broadcast over + the neighbor stream; ``gg`` is recomputed in registers (mirroring the + forward), and the channel axis is padded to ``NGP`` with masked loads. + + The residual gradient ``grad_h1`` folds the ``ng // H1`` residual copies. + When ``H1`` is a power of two the fold is the free reduction of the + ``(BLOCK_N, ng / H1, H1)`` view. When ``H1`` is padded (``H1P != H1``) + that view would split the copies at ``H1P`` rather than the true ``H1``, + so the doubling fold is instead evaluated directly as + ``grad_h2[:, :H1] + grad_h2[:, H1:2*H1]`` from the two upstream halves. + """ + node = tl.program_id(0) + rc = tl.arange(0, NGP) + cm = rc < NG + idt = tl.load(idt_ptr + rc, mask=cm, other=0.0) + gb = node * (4 * NG) + g0 = tl.load(gout_ptr + gb + 0 * NG + rc, mask=cm, other=0.0) + g1 = tl.load(gout_ptr + gb + 1 * NG + rc, mask=cm, other=0.0) + g2 = tl.load(gout_ptr + gb + 2 * NG + rc, mask=cm, other=0.0) + g3 = tl.load(gout_ptr + gb + 3 * NG + rc, mask=cm, other=0.0) + # Doubling with padded H1 needs the moment gradient at both residual + # halves; hoist the per-node loads at columns ``rj`` and ``rj + H1``. + if RESNET_MULT == 2 and H1P != H1: + rj = tl.arange(0, H1P) + hm = rj < H1 + g0lo = tl.load(gout_ptr + gb + 0 * NG + rj, mask=hm, other=0.0) + g1lo = tl.load(gout_ptr + gb + 1 * NG + rj, mask=hm, other=0.0) + g2lo = tl.load(gout_ptr + gb + 2 * NG + rj, mask=hm, other=0.0) + g3lo = tl.load(gout_ptr + gb + 3 * NG + rj, mask=hm, other=0.0) + g0hi = tl.load(gout_ptr + gb + 0 * NG + rj + H1, mask=hm, other=0.0) + g1hi = tl.load(gout_ptr + gb + 1 * NG + rj + H1, mask=hm, other=0.0) + g2hi = tl.load(gout_ptr + gb + 2 * NG + rj + H1, mask=hm, other=0.0) + g3hi = tl.load(gout_ptr + gb + 3 * NG + rj + H1, mask=hm, other=0.0) + for n0 in range(0, NNEI, BN): + offs = n0 + tl.arange(0, BN) + nmask = offs < NNEI + m = (nmask[:, None] & cm[None, :]) if NGP != NG else nmask[:, None] + e = node * NNEI + offs + ec = e[:, None] * NG + rc[None, :] + z2 = tl.load(z2_ptr + ec, mask=m, other=0.0) + base = e * 4 + s = tl.load(rr_ptr + base + 0, mask=nmask, other=0.0) + rx = tl.load(rr_ptr + base + 1, mask=nmask, other=0.0) + ry = tl.load(rr_ptr + base + 2, mask=nmask, other=0.0) + rz = tl.load(rr_ptr + base + 3, mask=nmask, other=0.0) + a, ad = activation_grad(z2, ACT) + h2 = a * idt[None, :] + if RESNET_MULT > 0: + h1d = tl.load( + h1_ptr + e[:, None] * H1 + (rc % H1)[None, :], mask=m, other=0.0 + ) + h2 = h2 + h1d + if GATED: + idx = tl.load(idx_ptr + e, mask=nmask, other=0) + ggt = tl.load( + tt_ptr + idx[:, None] * NG + rc[None, :], mask=m, other=0.0 + ) + sw = tl.load(sw_ptr + e, mask=nmask, other=0.0) + fac = 1.0 + ggt * sw[:, None] # d gg / d h2 + else: + fac = 1.0 + gg = h2 * fac + dgg = ( + s[:, None] * g0[None, :] + + rx[:, None] * g1[None, :] + + ry[:, None] * g2[None, :] + + rz[:, None] * g3[None, :] + ) + grad_h2 = dgg * fac + tl.store( + dz2_ptr + ec, + grad_h2 * idt[None, :] * ad, + mask=m, + ) + if RESNET_MULT == 1: + # Identity residual: ``grad_h1`` equals the moment gradient. + tl.store(dh1_ptr + ec, grad_h2, mask=m) + elif RESNET_MULT == 2 and H1P == H1: + grad_h1 = tl.sum(tl.reshape(grad_h2, (BN, 2, H1)), axis=1) + tl.store( + dh1_ptr + e[:, None] * H1 + tl.arange(0, H1)[None, :], + grad_h1, + mask=nmask[:, None], + ) + elif RESNET_MULT == 2: + # Padded doubling: recompute the two residual halves directly. + hmask = nmask[:, None] & (rj < H1)[None, :] + dgg_lo = ( + s[:, None] * g0lo[None, :] + + rx[:, None] * g1lo[None, :] + + ry[:, None] * g2lo[None, :] + + rz[:, None] * g3lo[None, :] + ) + dgg_hi = ( + s[:, None] * g0hi[None, :] + + rx[:, None] * g1hi[None, :] + + ry[:, None] * g2hi[None, :] + + rz[:, None] * g3hi[None, :] + ) + if GATED: + ggt_lo = tl.load( + tt_ptr + idx[:, None] * NG + rj[None, :], mask=hmask, other=0.0 + ) + ggt_hi = tl.load( + tt_ptr + idx[:, None] * NG + rj[None, :] + H1, + mask=hmask, + other=0.0, + ) + grad_h1 = dgg_lo * (1.0 + ggt_lo * sw[:, None]) + dgg_hi * ( + 1.0 + ggt_hi * sw[:, None] + ) + else: + grad_h1 = dgg_lo + dgg_hi + tl.store(dh1_ptr + e[:, None] * H1 + rj[None, :], grad_h1, mask=hmask) + if GATED: + # sw enters only through the gate; concat has no sw gradient. + tl.store(dsw_ptr + e, tl.sum(dgg * h2 * ggt, axis=1), mask=nmask) + tl.store(drr_ptr + base + 0, tl.sum(gg * g0[None, :], axis=1), mask=nmask) + tl.store(drr_ptr + base + 1, tl.sum(gg * g1[None, :], axis=1), mask=nmask) + tl.store(drr_ptr + base + 2, tl.sum(gg * g2[None, :], axis=1), mask=nmask) + tl.store(drr_ptr + base + 3, tl.sum(gg * g3[None, :], axis=1), mask=nmask) + + +# ====================================================================== +# Dispatch, operator registration and public API +# ====================================================================== +def _use_triton(tensor: Tensor) -> bool: + return ( + TRITON_AVAILABLE + and tensor.is_cuda + and tensor.dtype in (torch.float32, torch.float64) + ) + + +def _se_conv_fwd_impl( + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + resnet_mult: int, + act: int, + gated: int, + block_n: int, + num_warps: int, +) -> Tensor: + if not _use_triton(z2): + return _se_conv_reference(z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated) + nfnl, nnei, ng = z2.shape + out = torch.empty((nfnl, 4, ng), dtype=z2.dtype, device=z2.device) + wrap_triton(_se_conv_fwd_kernel)[(nfnl,)]( + z2, + h1, + idt, + tt, + idx, + sw, + rr, + out, + nnei, + H1=h1.shape[-1], + NG=ng, + NGP=triton.next_power_of_2(ng), + RESNET_MULT=resnet_mult, + ACT=act, + GATED=gated, + BN=block_n, + num_warps=num_warps, + ) + return out + + +def _se_conv_bwd_impl( + grad_xyz: Tensor, + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + resnet_mult: int, + act: int, + gated: int, + block_n: int, + num_warps: int, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + if not _use_triton(z2): + return _se_conv_reference_backward( + grad_xyz, z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated + ) + nfnl, nnei, ng = z2.shape + dz2 = torch.empty_like(z2) + # ``h1`` enters the activation only through the residual; with no residual + # its moment gradient is zero (its GEMM gradient is formed outside the op). + dh1 = torch.empty_like(h1) if resnet_mult > 0 else torch.zeros_like(h1) + # Concat (gated == 0) uses no sw gate; the kernel skips the dsw store, so + # pre-zero it here. + dsw = torch.empty_like(sw) if gated else torch.zeros_like(sw) + drr = torch.empty_like(rr) + wrap_triton(_se_conv_bwd_kernel)[(nfnl,)]( + grad_xyz.contiguous(), + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dz2, + dh1, + dsw, + drr, + nnei, + H1=h1.shape[-1], + NG=ng, + NGP=triton.next_power_of_2(ng), + H1P=triton.next_power_of_2(h1.shape[-1]), + RESNET_MULT=resnet_mult, + ACT=act, + GATED=gated, + BN=block_n, + num_warps=num_warps, + ) + return dz2, dh1, dsw, drr + + +_se_conv_op = triton_op("dpa1_triton::se_conv", mutates_args=())(_se_conv_fwd_impl) +_se_conv_bwd_op = triton_op("dpa1_triton::se_conv_bwd", mutates_args=())( + _se_conv_bwd_impl +) + + +@_se_conv_op.register_fake +def _(z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated, block_n, num_warps): + return z2.new_empty((z2.shape[0], 4, z2.shape[2])) + + +@_se_conv_bwd_op.register_fake +def _( + grad_xyz, z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated, block_n, num_warps +): + return ( + torch.empty_like(z2), + torch.empty_like(h1), + torch.empty_like(sw), + torch.empty_like(rr), + ) + + +def _se_conv_setup_context(ctx, inputs, output): + z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated, block_n, num_warps = inputs + ctx.save_for_backward(z2, h1, idt, tt, idx, sw, rr) + ctx.resnet_mult = resnet_mult + ctx.act = act + ctx.gated = gated + ctx.block_n = block_n + ctx.num_warps = num_warps + + +def _se_conv_backward(ctx, grad_xyz): + z2, h1, idt, tt, idx, sw, rr = ctx.saved_tensors + grad_z2, grad_h1, grad_sw, grad_rr = _se_conv_bwd_op( + grad_xyz.contiguous(), + z2, + h1, + idt, + tt, + idx, + sw, + rr, + ctx.resnet_mult, + ctx.act, + ctx.gated, + ctx.block_n, + ctx.num_warps, + ) + # ``idt`` / ``tt`` / ``idx`` do not depend on coordinates (no force grad). + return ( + grad_z2, + grad_h1, + None, + None, + None, + grad_sw, + grad_rr, + None, + None, + None, + None, + None, + ) + + +_se_conv_op.register_autograd(_se_conv_backward, setup_context=_se_conv_setup_context) + + +def concat_gate_placeholders(z2: Tensor, ng: int) -> tuple[Tensor, Tensor, Tensor]: + """Unused ``(tt, idx, sw)`` gate inputs for the concat (``gated == 0``) call. + + Concat carries the type feature through the embedding input, so the strip + type-pair gate is inactive and :func:`se_conv` skips these tensors. They are + still required to fill the operator signature; minimal (single-element) + placeholders suffice and keep the traced graph lean. + + Parameters + ---------- + z2 : Tensor + A tensor sharing the target device and floating dtype. + ng : int + Embedding channel width (the ``tt`` table's column count). + + Returns + ------- + tuple[Tensor, Tensor, Tensor] + The placeholder ``(tt, idx, sw)``: ``tt`` (1, ng), ``idx`` (1,) int64, + ``sw`` (1, 1). + """ + tt = z2.new_zeros(1, ng) + idx = torch.zeros(1, dtype=torch.long, device=z2.device) + sw = z2.new_ones(1, 1) + return tt, idx, sw + + +def se_conv( + z2: Tensor, + h1: Tensor, + idt: Tensor, + tt: Tensor, + idx: Tensor, + sw: Tensor, + rr: Tensor, + resnet_mult: int, + act: int, + gated: int, +) -> Tensor: + """Fused ``se_atten`` environment convolution (attn-free, strip or concat). + + Parameters + ---------- + z2 : Tensor + Last embedding pre-activation ``h1 @ W2 + b2`` with shape + (N, nnei, ng), where ``N = nframes * nloc``. + h1 : Tensor + Penultimate embedding activation with shape (N, nnei, H1); read only + when ``resnet_mult > 0``, where ``ng == resnet_mult * H1``. + idt : Tensor + Last-layer per-channel timestep with shape (ng,); all ones when the + network has no ``resnet_dt``. + tt : Tensor + Type-pair embedding table with shape (P, ng). + idx : Tensor + Per-edge row index into ``tt`` with shape (N * nnei,), dtype int64. + sw : Tensor + Smooth radial cutoff with shape (N, nnei). + rr : Tensor + Environment matrix ``(s, s*x/r, s*y/r, s*z/r)`` with shape (N, nnei, 4). + resnet_mult : int + Residual structure of the last layer: ``2`` (width doubling), ``1`` + (identity), or ``0`` (no residual). + act : int + Last-layer activation code from :data:`.activation.ACT_CODES`: ``0`` for + ``tanh``, ``1`` for ``silu``. + gated : int + Tebd-input mode: ``1`` (strip) applies the type-pair gate + ``gg = h2 * (1 + tt[idx] * sw)``; ``0`` (concat) skips it (``gg = h2``) + and ignores ``tt`` / ``idx`` / ``sw``, since concat feeds the type + feature through the embedding input. + + Returns + ------- + Tensor + Unnormalized moments ``xyz`` with shape (N, 4, ng), equal to + ``rr^T @ gg``. The ``1 / nnei`` normalization is applied by the + descriptor after this operator. + + Notes + ----- + The launch configuration is resolved from the active ``DP_TRITON_INFER`` + level via :func:`.tile_configs.resolve_conv_config`. The operator composes + under ``make_fx`` / ``torch.compile`` and exposes an inference force + gradient through the registered backward. + """ + block_n, num_warps = resolve_conv_config( + int(z2.shape[-1]), int(h1.shape[-1]), triton_infer_level() + ) + return _se_conv_op( + z2, h1, idt, tt, idx, sw, rr, resnet_mult, act, gated, block_n, num_warps + ) + + +def se_atten_conv( + embedding_net, + ss: Tensor, + tt: Tensor | None, + idx: Tensor | None, + sw: Tensor | None, + rr: Tensor, + gated: int, +) -> Tensor: + """Fuse the embedding net's final layer, type gate and moment accumulation. + + The embedding net is evaluated through its penultimate layer on the eager + (cuBLAS) path -- fp32 GEMMs which Triton cannot beat without tensor cores -- + and the final layer's activation (its timestep and residual) is folded into + :func:`se_conv` together with the type-pair gate (strip only), smooth cutoff + and the environment moment reduction. + + Parameters + ---------- + embedding_net : EmbeddingNet + The strip-mode radial embedding network. Any channel width is supported; + widths that are not powers of two are handled by masked padding inside + :func:`se_conv`. The last layer's activation must be one of + :data:`.activation.ACT_CODES` (``tanh`` or ``silu``); the caller verifies + this before routing here. + ss : Tensor + Embedding-net input with shape (N, nnei, d_in): the radial channel + (d_in == 1) for strip (``gated``), or the radial-plus-type-embedding + concatenation for concat (``not gated``). + tt : Tensor or None + Type-pair embedding table with shape (P, ng); the strip gate table. + Ignored (pass ``None``) for concat, whose type feature is already in + ``ss``. + idx : Tensor or None + Per-edge row index into ``tt`` with shape (N * nnei,); ``None`` for + concat. + sw : Tensor or None + Smooth radial cutoff with shape (N, nnei); ``None`` for concat. + rr : Tensor + Environment matrix with shape (N, nnei, 4). + gated : int + ``1`` (strip) applies the type-pair gate; ``0`` (concat) skips it. + + Returns + ------- + Tensor + Unnormalized moments ``xyz`` with shape (N, 4, ng). + """ + *head, last = embedding_net.layers + h = ss + for layer in head: + h = layer(h) + z2 = embed_last_gemm(h, last.matrix, last.bias) + h1_dim, ng = last.matrix.shape + if last.resnet and ng == 2 * h1_dim: + resnet_mult = 2 + elif last.resnet and ng == h1_dim: + resnet_mult = 1 + else: + resnet_mult = 0 + idt = last.idt if last.idt is not None else z2.new_ones(ng) + act = ACT_CODES[last.activate_name] + if not gated: + tt, idx, sw = concat_gate_placeholders(z2, ng) + return se_conv( + z2.contiguous(), h.contiguous(), idt, tt, idx, sw, rr, resnet_mult, act, gated + ) + + +# ====================================================================== +# Freeze-time launch-configuration autotuning (DP_TRITON_INFER >= 2) +# ====================================================================== +def _autotune_conv(model: torch.nn.Module, level: int, device: torch.device) -> None: + """Sweep the fused-convolution launch table for a model about to be frozen. + + Collects the ``(ng, H1)`` shape key of every eligible ``se_atten`` + descriptor in ``model`` and sweeps the keys the built-in / freeze-time + tables do not yet cover on the target GPU, registering the winners so the + ``resolve_conv_config`` lookups made while tracing bake tuned launches into + the exported ``.pt2``. Keys already covered cost nothing. + """ + from deepmd.kernels.triton.dpa1.sweep_tile_configs import ( + sweep, + ) + + device_name = torch.cuda.get_device_name(device) + keys: set[tuple[int, int]] = set() + for module in model.modules(): + eligible = getattr(module, "_fused_eligible", None) + if not (callable(eligible) and eligible("triton")): + continue + weight = module.se_atten.embeddings[0].layers[-1].w + keys.add((int(weight.shape[1]), int(weight.shape[0]))) + tuned: dict[tuple[int, int], tuple[int, int]] = {} + for ng, h1 in sorted(keys): + # The sweep needs a residual last layer (ng in {H1, 2*H1}); other shapes + # keep the default. Skip keys the tables already cover. + if has_conv_config(ng, h1) or ng not in (h1, 2 * h1): + continue + config = sweep(ng, h1, device=device) + register_conv_config(device_name, ng, h1, config) + tuned[(ng, h1)] = config + if tuned: + log.info("DPA1 se_conv: tuned launch configs %s on %s.", tuned, device_name) + else: + log.info( + "DPA1 se_conv: launch table already covers this checkpoint's shapes " + "on %s; no tuning needed.", + device_name, + ) + + +if TRITON_AVAILABLE: + register_autotuner("dpa1_se_conv", _autotune_conv) diff --git a/deepmd/kernels/triton/dpa1/sweep_tile_configs.py b/deepmd/kernels/triton/dpa1/sweep_tile_configs.py new file mode 100644 index 0000000000..51e81bb2c9 --- /dev/null +++ b/deepmd/kernels/triton/dpa1/sweep_tile_configs.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +# ruff: noqa: T201 +r"""Sweep the launch configuration of a DPA1 fused environment convolution. + +Both fused kernels have a two-parameter launch configuration resolved per +``(ng, H1)`` by :mod:`.tile_configs`: ``se_conv`` (node-parallel) is keyed by +the per-neighbor block width ``BLOCK_N`` and a warp count; ``edge_conv`` +(edge-parallel) by the per-block edge count ``BLOCK_E`` and a warp count. This +module measures the candidate configurations for one channel width on synthetic +tensors at a production count, validating each against the eager reference +before timing, and reports the fastest forward+backward configuration that +stays spill-free. + +The forward is memory-bound and tolerant of the configuration; the backward +holds a ``(BLOCK, ng)`` block live and collapses if it spills, so the sweep +scores forward+backward jointly. + +Usage +----- +:: + + python -m deepmd.kernels.triton.dpa1.sweep_tile_configs \\ + --kind {conv,edge} --ng NG --h1 H1 [--device cuda:0] + +The printed ``(ng, h1): (BLOCK, num_warps)`` line is appended, under the device +name from ``torch.cuda.get_device_name``, to the relevant built-in table in +:mod:`.tile_configs` (``_CONV_BUILTIN`` / ``_EDGE_BUILTIN``). + +Regeneration note +----------------- +Any change to a kernel body invalidates its existing table entries; rerun the +sweep for every covered ``(ng, H1)`` and refresh the table. +""" + +from __future__ import ( + annotations, +) + +import argparse +import itertools +from typing import ( + TYPE_CHECKING, +) + +import torch +from torch import ( + Tensor, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + +from deepmd.kernels.triton.dpa1.edge_conv import ( + _edge_conv_bwd_impl, + _edge_conv_fwd_impl, + _edge_conv_reference, +) +from deepmd.kernels.triton.dpa1.se_conv import ( + _se_conv_bwd_impl, + _se_conv_fwd_impl, + _se_conv_reference, +) +from deepmd.kernels.triton.dpa1.tile_configs import ( + EDGE_DEFAULT_CONFIG, +) + +# BLOCK_N candidates are powers of two; small values protect the backward +# register footprint, large values amortize the per-node loop overhead. +_BLOCK_N_CANDIDATES = (16, 32, 64, 128) +# BLOCK_E candidates (edges per program) are powers of two; small values keep +# the per-program register tile small, large values amortize launch overhead. +_BLOCK_E_CANDIDATES = (1, 2, 4, 8, 16, 32) +_WARP_CANDIDATES = (2, 4, 8) +_REL_TOL = 1e-5 +# An edge_conv candidate is only recorded when it beats the universal default by +# more than ``1 - _EDGE_WIN_RATIO`` (5%); otherwise the default is kept, so a +# level-2 launch never regresses below level 1 across edge counts. +_EDGE_WIN_RATIO = 0.95 + + +def _make_inputs( + nodes: int, nnei: int, ng: int, h1: int, device: torch.device +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + p = 4096 + z2 = torch.randn(nodes, nnei, ng, dtype=torch.float32, device=device) + h1t = torch.randn(nodes, nnei, h1, dtype=torch.float32, device=device) + idt = torch.ones(ng, dtype=torch.float32, device=device) + tt = torch.randn(p, ng, dtype=torch.float32, device=device) * 0.3 + idx = torch.randint(0, p, (nodes * nnei,), dtype=torch.int64, device=device) + sw = torch.rand(nodes, nnei, dtype=torch.float32, device=device) + rr = torch.randn(nodes, nnei, 4, dtype=torch.float32, device=device) + return z2, h1t, idt, tt, idx, sw, rr + + +def _bench(fn: Callable[[], object], iters: int = 40, warmup: int = 15) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def sweep( + ng: int, + h1: int, + nnei: int = 181, + nodes: int = 4096, + device: torch.device | None = None, +) -> tuple[int, int]: + """Measure and return the fastest spill-free ``(BLOCK_N, num_warps)``. + + Parameters + ---------- + ng : int + Embedding channel width. + h1 : int + Penultimate embedding width; ``ng == 2 * h1`` is required. + nnei : int + Neighbor count used to size the synthetic input. + nodes : int + Node count used to size the synthetic input. + device : torch.device, optional + CUDA device; defaults to the current device. + + Returns + ------- + tuple[int, int] + The winning ``(BLOCK_N, num_warps)``. + """ + if ng not in (h1, 2 * h1): + raise ValueError( + "se_conv sweep requires a residual last layer (ng in {h1, 2*h1})" + ) + resnet_mult = ng // h1 + device = device or torch.device("cuda") + torch.backends.cuda.matmul.allow_tf32 = False + # The launch configuration is memory/register bound and keyed by (ng, H1) + # only; the activation adds a few cheap elementwise ops and does not shift + # the optimum, so the sweep times the ``tanh`` path (act = 0). + act = 0 + # The launch configuration is keyed by (ng, H1) and is independent of the + # tebd-input mode; the strip gate (gated = 1) is the register-heaviest case, + # so its optimum is a safe upper bound for concat (gated = 0). + gated = 1 + z2, h1t, idt, tt, idx, sw, rr = _make_inputs(nodes, nnei, ng, h1, device) + ref = _se_conv_reference(z2, h1t, idt, tt, idx, sw, rr, resnet_mult, act, gated) + gout = torch.randn_like(ref) + + def fwd_bwd(bn: int, nw: int) -> None: + _se_conv_fwd_impl( + z2, h1t, idt, tt, idx, sw, rr, resnet_mult, act, gated, bn, nw + ) + _se_conv_bwd_impl( + gout, z2, h1t, idt, tt, idx, sw, rr, resnet_mult, act, gated, bn, nw + ) + + results: list[tuple[float, int, int]] = [] + for bn, nw in itertools.product(_BLOCK_N_CANDIDATES, _WARP_CANDIDATES): + try: + out = _se_conv_fwd_impl( + z2, h1t, idt, tt, idx, sw, rr, resnet_mult, act, gated, bn, nw + ) + rel = (out - ref).abs().max().item() / ref.abs().max().item() + if rel > _REL_TOL: + continue + t = _bench(lambda: fwd_bwd(bn, nw)) + except Exception: + continue + results.append((t, bn, nw)) + print(f" BLOCK_N={bn:<4} num_warps={nw}: fwd+bwd {t * 1e3:8.1f} us") + if not results: + raise RuntimeError("no valid configuration found") + best_t, best_bn, best_nw = min(results, key=lambda r: r[0]) + return best_bn, best_nw + + +def _make_edge_inputs( + n_edge: int, n_node: int, ng: int, h1: int, device: torch.device +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + p = 4096 + z2 = torch.randn(n_edge, ng, dtype=torch.float32, device=device) + h1t = torch.randn(n_edge, h1, dtype=torch.float32, device=device) + idt = torch.ones(ng, dtype=torch.float32, device=device) + tt = torch.randn(p, ng, dtype=torch.float32, device=device) * 0.3 + idx = torch.randint(0, p, (n_edge,), dtype=torch.int64, device=device) + sw = torch.rand(n_edge, dtype=torch.float32, device=device) + rr = torch.randn(n_edge, 4, dtype=torch.float32, device=device) + dst = torch.randint(0, n_node, (n_edge,), dtype=torch.int64, device=device) + edge_mask = (torch.rand(n_edge, dtype=torch.float32, device=device) > 0.05).to( + dtype=torch.float32 + ) + return z2, h1t, idt, tt, idx, sw, rr, dst, edge_mask + + +def sweep_edge( + ng: int, + h1: int, + n_edge: int = 524288, + n_node: int = 4096, + device: torch.device | None = None, +) -> tuple[int, int]: + """Measure and return the fastest spill-free ``edge_conv`` ``(BLOCK_E, num_warps)``. + + Unlike ``se_conv`` (whose optimum is edge-count-insensitive), the + edge-parallel launch's optimum drifts mildly with the edge count through + occupancy, so the sweep runs at a large, throughput-bound count where the + ratios stabilize, and only returns a tuned configuration when it beats the + universal default (:data:`.tile_configs.EDGE_DEFAULT_CONFIG`) by more than + ``1 - _EDGE_WIN_RATIO``; otherwise it returns the default so a subsequent + level-2 launch never regresses below level 1. + + Parameters + ---------- + ng : int + Embedding channel width. + h1 : int + Penultimate embedding width; ``ng in {h1, 2 * h1}`` for a residual layer. + n_edge : int + Edge count used to size the synthetic input. + n_node : int + Node count (segment count) used to size the scatter target. + device : torch.device, optional + CUDA device; defaults to the current device. + + Returns + ------- + tuple[int, int] + The winning ``(BLOCK_E, num_warps)``, or the universal default when no + candidate is a clear win. + """ + resnet_mult = ng // h1 if ng in (h1, 2 * h1) else 0 + device = device or torch.device("cuda") + torch.backends.cuda.matmul.allow_tf32 = False + act = 0 + # The launch configuration is keyed by (ng, H1) and is independent of the + # tebd-input mode; the strip gate (gated = 1) is the register-heaviest case, + # so its optimum is a safe upper bound for concat (gated = 0). + gated = 1 + # Shared leading arguments of the three edge_conv entry points, ordered to + # match their signatures so each call splats this tuple and appends only its + # launch configuration. + edge_args = ( + *_make_edge_inputs(n_edge, n_node, ng, h1, device), + n_node, + resnet_mult, + act, + gated, + ) + ref = _edge_conv_reference(*edge_args) + gout = torch.randn_like(ref) + + def fwd_bwd(be: int, nw: int) -> None: + _edge_conv_fwd_impl(*edge_args, be, nw) + _edge_conv_bwd_impl(gout, *edge_args, be, nw) + + results: dict[tuple[int, int], float] = {} + for be, nw in itertools.product(_BLOCK_E_CANDIDATES, _WARP_CANDIDATES): + try: + out = _edge_conv_fwd_impl(*edge_args, be, nw) + rel = (out - ref).abs().max().item() / ref.abs().max().item() + if rel > _REL_TOL: + continue + t = _bench(lambda: fwd_bwd(be, nw)) + except Exception: + continue + results[(be, nw)] = t + print(f" BLOCK_E={be:<4} num_warps={nw}: fwd+bwd {t * 1e3:8.1f} us") + if not results: + raise RuntimeError("no valid configuration found") + best = min(results, key=results.get) + # Only accept a tuned win over the default; the default is always a candidate. + default_t = results.get(EDGE_DEFAULT_CONFIG) + if default_t is not None and results[best] >= default_t * _EDGE_WIN_RATIO: + return EDGE_DEFAULT_CONFIG + return best + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--kind", choices=("conv", "edge"), default="conv") + parser.add_argument("--ng", type=int, required=True) + parser.add_argument("--h1", type=int, required=True) + parser.add_argument("--nnei", type=int, default=181) + parser.add_argument("--nodes", type=int, default=4096) + parser.add_argument("--device", type=str, default="cuda:0") + args = parser.parse_args() + device = torch.device(args.device) + torch.cuda.set_device(device) + if args.kind == "edge": + block, nw = sweep_edge(args.ng, args.h1, device=device) + else: + block, nw = sweep(args.ng, args.h1, args.nnei, args.nodes, device) + print(f'\n"{torch.cuda.get_device_name()}": {{') + print(f" ({args.ng}, {args.h1}): ({block}, {nw}),") + print("}") + + +if __name__ == "__main__": + main() diff --git a/deepmd/kernels/triton/dpa1/tile_configs.py b/deepmd/kernels/triton/dpa1/tile_configs.py new file mode 100644 index 0000000000..0218446231 --- /dev/null +++ b/deepmd/kernels/triton/dpa1/tile_configs.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Launch-configuration resolution for the DPA1 fused environment convolutions. + +Two memory-bound kernels are configured here, each reducing to a block width +and a warp count keyed by the channel width ``ng`` and the resnet width ``H1``: + +- ``se_conv`` (strip / dense, node-parallel): one program owns a node and + streams its neighbors, so the launch is ``(BLOCK_N, num_warps)`` -- neighbors + per block. ``BLOCK_N`` bounds the live ``(BLOCK_N, channels)`` register + footprint of the backward pass; oversized blocks spill and collapse + throughput, so the universal default is kept small. +- ``edge_conv`` (concat / graph, edge-parallel): one program owns a block of + edges and scatters them into their center nodes, so the launch is + ``(BLOCK_E, num_warps)`` -- edges per block. ``BLOCK_E`` bounds the live + ``(BLOCK_E, channels)`` register footprint of both passes. + +The optimum depends on ``(ng, H1)`` (the register footprint) but is insensitive +to the neighbor / edge count, which only sets the loop trip count or the grid +size. Table keys are therefore ``(ng, H1)``. + +Level policy (see :func:`deepmd.kernels.utils.triton_infer_level`): + +- Level ``1`` always returns the family default, a single shape-independent + configuration that never spills. +- Level ``>= 2`` consults the per-GPU table and falls back to the default on any + miss, so an unswept device or channel width can only match or beat level 1, + never regress below it. + +New devices are added by running :mod:`.sweep_tile_configs` and appending a +block to the relevant built-in table under the device name reported by +``torch.cuda.get_device_name``. +""" + +from __future__ import ( + annotations, +) + +import torch + +# (block, num_warps): ``block`` is ``BLOCK_N`` for se_conv and ``BLOCK_E`` for +# edge_conv. A small block keeps the backward register footprint within budget. +Config = tuple[int, int] + +# se_conv (node-parallel) universal default. +DEFAULT_CONFIG: Config = (16, 4) +# edge_conv (edge-parallel) universal default. +EDGE_DEFAULT_CONFIG: Config = (8, 4) + +# Per-GPU built-in tables keyed by (ng, H1). Values are the fastest spill-free +# forward+backward configuration produced by the fp32 sweep in +# :mod:`.sweep_tile_configs` for that channel width. +_CONV_BUILTIN: dict[str, dict[tuple[int, int], Config]] = { + "NVIDIA H20": { + (32, 16): (32, 2), + (64, 32): (16, 2), + (128, 64): (16, 2), + (256, 128): (16, 4), + (100, 50): (16, 2), + (200, 100): (16, 4), + }, +} +_EDGE_BUILTIN: dict[str, dict[tuple[int, int], Config]] = { + "NVIDIA H20": { + (32, 16): (8, 4), + (64, 32): (8, 2), + (128, 64): (8, 2), + (256, 128): (2, 2), + (100, 50): (8, 2), + (200, 100): (1, 2), + (32, 32): (8, 4), + (64, 64): (8, 2), + (128, 128): (1, 2), + }, +} + +# Per-GPU configs swept at freeze time by :mod:`deepmd.kernels.autotune` for +# shape keys the built-in tables do not cover. Process-local: the freeze traces +# on the target GPU, so these are baked into the exported ``.pt2``; they never +# persist across processes. Same schema as the built-in tables. +_CONV_RUNTIME: dict[str, dict[tuple[int, int], Config]] = {} +_EDGE_RUNTIME: dict[str, dict[tuple[int, int], Config]] = {} + + +def _register( + runtime: dict[str, dict[tuple[int, int], Config]], + device_name: str, + ng: int, + h1: int, + config: Config, +) -> None: + runtime.setdefault(device_name, {})[(int(ng), int(h1))] = config + + +def _covered( + builtin: dict[str, dict[tuple[int, int], Config]], + runtime: dict[str, dict[tuple[int, int], Config]], + ng: int, + h1: int, +) -> bool: + if not torch.cuda.is_available(): + return False + name = torch.cuda.get_device_name() + key = (int(ng), int(h1)) + return key in runtime.get(name, {}) or key in builtin.get(name, {}) + + +def _resolve( + builtin: dict[str, dict[tuple[int, int], Config]], + runtime: dict[str, dict[tuple[int, int], Config]], + default: Config, + ng: int, + h1: int, + level: int, +) -> Config: + if level < 2 or not torch.cuda.is_available(): + return default + name = torch.cuda.get_device_name() + key = (int(ng), int(h1)) + runtime_dev = runtime.get(name, {}) + if key in runtime_dev: + return runtime_dev[key] + return builtin.get(name, {}).get(key, default) + + +# --- se_conv (node-parallel) ------------------------------------------------ +def register_conv_config(device_name: str, ng: int, h1: int, config: Config) -> None: + """Register a freshly swept ``se_conv`` launch for ``(ng, h1)``. + + Used by the freeze-time autotuner so a subsequent :func:`resolve_conv_config` + (made while tracing) bakes the tuned launch into the exported artifact. + """ + _register(_CONV_RUNTIME, device_name, ng, h1, config) + + +def has_conv_config(ng: int, h1: int) -> bool: + """Whether a tuned ``se_conv`` entry (built-in or freeze-time) covers ``(ng, h1)``.""" + return _covered(_CONV_BUILTIN, _CONV_RUNTIME, ng, h1) + + +def resolve_conv_config(ng: int, h1: int, level: int) -> Config: + """Resolve the ``(BLOCK_N, num_warps)`` for a fused ``se_conv`` launch. + + Level 1 forces the universal default; level ``>= 2`` consults the + freeze-time and per-GPU tables with fallback. ``BLOCK_N`` is a power of two. + """ + return _resolve(_CONV_BUILTIN, _CONV_RUNTIME, DEFAULT_CONFIG, ng, h1, level) + + +# --- edge_conv (edge-parallel) ---------------------------------------------- +def register_edge_config(device_name: str, ng: int, h1: int, config: Config) -> None: + """Register a freshly swept ``edge_conv`` launch for ``(ng, h1)``.""" + _register(_EDGE_RUNTIME, device_name, ng, h1, config) + + +def has_edge_config(ng: int, h1: int) -> bool: + """Whether a tuned ``edge_conv`` entry (built-in or freeze-time) covers ``(ng, h1)``.""" + return _covered(_EDGE_BUILTIN, _EDGE_RUNTIME, ng, h1) + + +def resolve_edge_config(ng: int, h1: int, level: int) -> Config: + """Resolve the ``(BLOCK_E, num_warps)`` for a fused ``edge_conv`` launch. + + Level 1 forces the universal default; level ``>= 2`` consults the + freeze-time and per-GPU tables with fallback. ``BLOCK_E`` is a power of two. + """ + return _resolve(_EDGE_BUILTIN, _EDGE_RUNTIME, EDGE_DEFAULT_CONFIG, ng, h1, level) diff --git a/deepmd/kernels/triton/env_mat.py b/deepmd/kernels/triton/env_mat.py new file mode 100644 index 0000000000..bad288b8ac --- /dev/null +++ b/deepmd/kernels/triton/env_mat.py @@ -0,0 +1,1116 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +# pyright: reportMissingImports=false +# ruff: noqa: ANN001, ANN202 +"""Universal fused smooth environment-matrix kernel (``DP_TRITON_INFER >= 1``). + +The smooth environment matrix is the shared front end of every ``se``-family +descriptor (``se_a``, ``se_r``, ``se_t``, ``se_atten`` / DPA1, and the DPA2/DPA3 +representation blocks): from the neighbor list it forms, per center ``i`` and +neighbor slot ``n`` (with relative vector ``d = r_j - r_i`` and distance +``L = |d|``), the normalized quartet + + env[0] = (s(L) / (L + p) - avg[0]) / std[0] + env[1..3] = (s(L) * d / (L + p)^2 - avg[1..3]) / std[1..3] + +where ``s`` is the smooth switch (quintic ``compute_smooth_weight`` or the +exponential ``compute_exp_sw``), ``p`` the division protection, and ``avg``/ +``std`` the per-type running statistics. It also returns the relative vector +``diff = d`` and the switch value ``sw = s(L)``; invalid / padding slots +(``nlist < 0``) carry a zero switch, a zero ``diff``, and hence ``env = -avg/std`` +(the switch multiplies the raw quartet before the affine normalization). + +The eager path expresses this as a chain of roughly ten memory-bound tensor ops +whose autograd backward -- the force path -- expands into as many gather / norm / +switch / division gradient kernels; on a 4k-atom cell that backward dominates the +front end (~0.75 ms). This module fuses the forward into a single node-parallel +kernel and replaces the autograd backward with a closed-form one: one kernel +evaluates the ``d env / d d`` Jacobian (contracted against the upstream grads of +all three outputs) into a per-edge ``d`` gradient, and a single scatter-add sends +it to the neighbor (``+``) and center (``-``) coordinates. + +The analytic Jacobian per slot, with ``q = L + p``, ``W = s(L)`` (masked), +``W' = s'(L)`` (masked), ``g_raw = g_env / std``, and ``G = ``:: + + g_d = (d / L) * [ g_raw[0] (W'/q - W/q^2) + G (W'/q^2 - 2 W/q^3) ] + + (W / q^2) * g_raw[1..3] + + g_diff + + g_sw * W' * (d / L) + +The kernel is inference-only (``register_autograd`` provides the first-order +backward used for forces; higher-order / training differentiation keeps the eager +path) and is registered as a ``triton_op`` so it is captured as a single opaque +node under ``make_fx`` / ``torch.export`` (the ``pt_expt`` backend). Off CUDA or +with Triton unavailable it transparently falls back to the validated eager +reference, so it is a drop-in for the descriptors' ``prod_env_mat`` front end. +""" + +from __future__ import ( + annotations, +) + +import torch +from torch import ( + Tensor, +) +from torch.library import ( + triton_op, + wrap_triton, +) + +from deepmd.dpmodel.utils.safe_gradient import ( + safe_for_vector_norm, +) +from deepmd.kernels.utils import ( + triton_infer_level, +) + +# Constant of the exponential switch ``compute_exp_sw`` (``a = C / rmin``). +_EXP_SWITCH_C = 20.0 + +try: + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: + TRITON_AVAILABLE = False + +__all__ = [ + "TRITON_AVAILABLE", + "edge_env_mat", + "env_mat", +] + + +if TRITON_AVAILABLE: + + @triton.jit + def _switch(L, rmin, rmax, sw_c, USE_EXP: tl.constexpr): + """Smooth switch ``s(L)`` and its derivative ``s'(L)`` (clamped window). + + Matches ``compute_smooth_weight`` (quintic) / ``compute_exp_sw`` + (double exponential); the derivative is zero outside the smooth window, + consistent with the ``clamp`` in the eager definitions. + """ + if USE_EXP: + a = sw_c / rmin + inside = (L > 0.0) & (L < rmax) + lc = tl.minimum(tl.maximum(L, 0.0), rmax) + xarg = a * (lc - rmin) + e = tl.exp(xarg) + w = tl.exp(-e) + # d(exp(-e))/dL = -a e w = -a exp(xarg - e); the fused exponent stays + # finite as e -> inf (xarg - e -> -inf), avoiding the 0*inf NaN that + # the factored ``-a e w`` produces once ``e`` overflows in fp32. + dw = tl.where(inside, -a * tl.exp(xarg - e), 0.0) + else: + inside = (L > rmin) & (L < rmax) + u = (L - rmin) / (rmax - rmin) + u = tl.minimum(tl.maximum(u, 0.0), 1.0) + u2 = u * u + w = u2 * u * (-6.0 * u2 + 15.0 * u - 10.0) + 1.0 + dw = tl.where( + inside, + (-30.0 * u2 * u2 + 60.0 * u2 * u - 30.0 * u2) / (rmax - rmin), + 0.0, + ) + return w, dw + + @triton.jit + def _env_mat_fwd_kernel( + coord_ptr, # (nf, nall, 3) + nlist_ptr, # (nf, nloc, nnei) int + atype_ptr, # (nf, nloc) int + mean_ptr, # (ntypes, nnei, C) + std_ptr, # (ntypes, nnei, C) + env_ptr, # (nf, nloc, nnei, C) + diff_ptr, # (nf, nloc, nnei, 3) + sw_ptr, # (nf, nloc, nnei) + rcut, + rcut_smth, + protection, + sw_c, + nloc, + nnei, + nall, + C: tl.constexpr, + RADIAL: tl.constexpr, + USE_EXP: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + """One program per center; a lane per neighbor slot (padded to BLOCK_N).""" + pid = tl.program_id(0) + f = pid // nloc + i = pid % nloc + n = tl.arange(0, BLOCK_N) + nmask = n < nnei + + ci = (f * nall + i) * 3 + cix = tl.load(coord_ptr + ci + 0) + ciy = tl.load(coord_ptr + ci + 1) + ciz = tl.load(coord_ptr + ci + 2) + ti = tl.load(atype_ptr + f * nloc + i) + + jj = tl.load(nlist_ptr + pid * nnei + n, mask=nmask, other=-1) + m = jj >= 0 + js = tl.where(m, jj, 0) + cj = (f * nall + js) * 3 + cjx = tl.load(coord_ptr + cj + 0, mask=nmask, other=0.0) + cjy = tl.load(coord_ptr + cj + 1, mask=nmask, other=0.0) + cjz = tl.load(coord_ptr + cj + 2, mask=nmask, other=0.0) + dx = cjx - cix + dy = cjy - ciy + dz = cjz - ciz + length = tl.sqrt(dx * dx + dy * dy + dz * dz) + # invalid slots take L=1 to keep q finite; the zero switch nulls the raw + # quartet and the diff, so their contribution vanishes regardless. + lsafe = tl.where(m, length, 1.0) + q = lsafe + protection + sw_raw, _ = _switch(lsafe, rcut_smth, rcut, sw_c, USE_EXP) + w = tl.where(m, sw_raw, 0.0) + inv_q = 1.0 / q + + base = pid * nnei * C + n * C + mb = ti * nnei * C + n * C + raw0 = w * inv_q + mean0 = tl.load(mean_ptr + mb, mask=nmask, other=0.0) + std0 = tl.load(std_ptr + mb, mask=nmask, other=1.0) + tl.store(env_ptr + base, (raw0 - mean0) / std0, mask=nmask) + tl.store(sw_ptr + pid * nnei + n, w, mask=nmask) + + db = pid * nnei * 3 + n * 3 + mf = tl.where(m, 1.0, 0.0) + tl.store(diff_ptr + db + 0, dx * mf, mask=nmask) + tl.store(diff_ptr + db + 1, dy * mf, mask=nmask) + tl.store(diff_ptr + db + 2, dz * mf, mask=nmask) + + if not RADIAL: + inv_q2 = inv_q * inv_q + wq2 = w * inv_q2 + for c in tl.static_range(3): + dc = tl.where(c == 0, dx, tl.where(c == 1, dy, dz)) + meanc = tl.load(mean_ptr + mb + (c + 1), mask=nmask, other=0.0) + stdc = tl.load(std_ptr + mb + (c + 1), mask=nmask, other=1.0) + tl.store( + env_ptr + base + (c + 1), (wq2 * dc - meanc) / stdc, mask=nmask + ) + + @triton.jit + def _env_mat_bwd_kernel( + coord_ptr, # (nf, nall, 3) + nlist_ptr, # (nf, nloc, nnei) + atype_ptr, # (nf, nloc) + std_ptr, # (ntypes, nnei, C) + genv_ptr, # (nf, nloc, nnei, C) + gdiff_ptr, # (nf, nloc, nnei, 3) + gsw_ptr, # (nf, nloc, nnei) + grad_ptr, # (nf, nall, 3) output: grad wrt coordinates (zero-initialized) + rcut, + rcut_smth, + protection, + sw_c, + nloc, + nnei, + nall, + C: tl.constexpr, + RADIAL: tl.constexpr, + USE_EXP: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + """Closed-form ``d out / d d`` scattered to the coordinate gradient. + + The per-slot ``d`` gradient goes to the neighbor coordinate (``+``, an + atomic add since neighbors are shared across centers) and, summed over + the slots, to the center coordinate (``-``); fusing the scatter here + avoids a chain of eager tensor ops (the eager backward's dominant + launch-overhead cost). + """ + pid = tl.program_id(0) + f = pid // nloc + i = pid % nloc + n = tl.arange(0, BLOCK_N) + nmask = n < nnei + + ci = (f * nall + i) * 3 + cix = tl.load(coord_ptr + ci + 0) + ciy = tl.load(coord_ptr + ci + 1) + ciz = tl.load(coord_ptr + ci + 2) + ti = tl.load(atype_ptr + f * nloc + i) + + jj = tl.load(nlist_ptr + pid * nnei + n, mask=nmask, other=-1) + m = jj >= 0 + js = tl.where(m, jj, 0) + cj = (f * nall + js) * 3 + dx = tl.load(coord_ptr + cj + 0, mask=nmask, other=0.0) - cix + dy = tl.load(coord_ptr + cj + 1, mask=nmask, other=0.0) - ciy + dz = tl.load(coord_ptr + cj + 2, mask=nmask, other=0.0) - ciz + length = tl.sqrt(dx * dx + dy * dy + dz * dz) + lsafe = tl.where(m, length, 1.0) + q = lsafe + protection + sw_raw, dsw_raw = _switch(lsafe, rcut_smth, rcut, sw_c, USE_EXP) + w = tl.where(m, sw_raw, 0.0) + wp = tl.where(m, dsw_raw, 0.0) + inv_q = 1.0 / q + inv_q2 = inv_q * inv_q + inv_q3 = inv_q2 * inv_q + inv_l = tl.where(m, 1.0 / lsafe, 0.0) + + mb = ti * nnei * C + n * C + gb = pid * nnei * C + n * C + std0 = tl.load(std_ptr + mb, mask=nmask, other=1.0) + g0 = tl.load(genv_ptr + gb, mask=nmask, other=0.0) / std0 + + gvx = tl.zeros((BLOCK_N,), dtype=g0.dtype) + gvy = tl.zeros((BLOCK_N,), dtype=g0.dtype) + gvz = tl.zeros((BLOCK_N,), dtype=g0.dtype) + gdot = tl.zeros((BLOCK_N,), dtype=g0.dtype) + wq2 = w * inv_q2 + if not RADIAL: + std1 = tl.load(std_ptr + mb + 1, mask=nmask, other=1.0) + std2 = tl.load(std_ptr + mb + 2, mask=nmask, other=1.0) + std3 = tl.load(std_ptr + mb + 3, mask=nmask, other=1.0) + gvx = tl.load(genv_ptr + gb + 1, mask=nmask, other=0.0) / std1 + gvy = tl.load(genv_ptr + gb + 2, mask=nmask, other=0.0) / std2 + gvz = tl.load(genv_ptr + gb + 3, mask=nmask, other=0.0) / std3 + gdot = gvx * dx + gvy * dy + gvz * dz + + gsw = tl.load(gsw_ptr + pid * nnei + n, mask=nmask, other=0.0) + # coef multiplies d/L: the g_env[0] + g_env[1..3] + g_sw radial terms. + coef = ( + g0 * (wp * inv_q - w * inv_q2) + + gdot * (wp * inv_q2 - 2.0 * w * inv_q3) + + gsw * wp + ) * inv_l + + db = pid * nnei * 3 + n * 3 + gdfx = tl.load(gdiff_ptr + db + 0, mask=nmask, other=0.0) + gdfy = tl.load(gdiff_ptr + db + 1, mask=nmask, other=0.0) + gdfz = tl.load(gdiff_ptr + db + 2, mask=nmask, other=0.0) + mf = tl.where(m, 1.0, 0.0) + gdx = (coef * dx + wq2 * gvx + gdfx) * mf + gdy = (coef * dy + wq2 * gvy + gdfy) * mf + gdz = (coef * dz + wq2 * gvz + gdfz) * mf + # scatter d(d = r_j - r_i): +grad to the neighbor row, -sum to the center. + smask = nmask & m + cj = (f * nall + js) * 3 + tl.atomic_add(grad_ptr + cj + 0, gdx, mask=smask) + tl.atomic_add(grad_ptr + cj + 1, gdy, mask=smask) + tl.atomic_add(grad_ptr + cj + 2, gdz, mask=smask) + tl.atomic_add(grad_ptr + ci + 0, -tl.sum(gdx, axis=0)) + tl.atomic_add(grad_ptr + ci + 1, -tl.sum(gdy, axis=0)) + tl.atomic_add(grad_ptr + ci + 2, -tl.sum(gdz, axis=0)) + + @triton.jit + def _edge_env_mat_fwd_kernel( + edge_vec_ptr, # (E, 3) + ctype_ptr, # (E,) int center-atom type + emask_ptr, # (E,) int valid-edge flag (1 valid, 0 padding) + mean_ptr, # (ntypes, 4) slot-independent + std_ptr, # (ntypes, 4) + env_ptr, # (E, 4) + sw_ptr, # (E,) smooth switch, zeroed on padding + rcut, + rcut_smth, + protection, + sw_c, + n_edge, + BLOCK_E: tl.constexpr, + ): + """Graph-native (slot-free) environment matrix; one lane per edge. + + The quintic switch is not masked in the env quartet: padding edges + (``valid == 0``) instead take ``length + 1`` so the switch decays past + the cutoff, matching the dense ``length + ~mask`` guard. The quartet + rows carry nonzero values on padding and are masked downstream by the + edge mask. The separately emitted switch ``sw`` is zeroed on padding, + mirroring the dense ``weight * mask`` (it feeds the strip type-pair gate). + """ + pid = tl.program_id(0) + e = pid * BLOCK_E + tl.arange(0, BLOCK_E) + m = e < n_edge + dx = tl.load(edge_vec_ptr + e * 3 + 0, mask=m, other=0.0) + dy = tl.load(edge_vec_ptr + e * 3 + 1, mask=m, other=0.0) + dz = tl.load(edge_vec_ptr + e * 3 + 2, mask=m, other=0.0) + ct = tl.load(ctype_ptr + e, mask=m, other=0) + valid = tl.load(emask_ptr + e, mask=m, other=0) + length = tl.sqrt(dx * dx + dy * dy + dz * dz) + length = length + tl.where(valid != 0, 0.0, 1.0) + denom = length + protection + sw, _ = _switch(length, rcut_smth, rcut, sw_c, 0) + inv_d = 1.0 / denom + inv_d2 = inv_d * inv_d + mb = ct * 4 + m0 = tl.load(mean_ptr + mb + 0, mask=m, other=0.0) + s0 = tl.load(std_ptr + mb + 0, mask=m, other=1.0) + tl.store(env_ptr + e * 4 + 0, (sw * inv_d - m0) / s0, mask=m) + tl.store(sw_ptr + e, tl.where(valid != 0, sw, 0.0), mask=m) + for c in tl.static_range(3): + dc = tl.where(c == 0, dx, tl.where(c == 1, dy, dz)) + mc = tl.load(mean_ptr + mb + (c + 1), mask=m, other=0.0) + sc = tl.load(std_ptr + mb + (c + 1), mask=m, other=1.0) + tl.store(env_ptr + e * 4 + (c + 1), (sw * dc * inv_d2 - mc) / sc, mask=m) + + @triton.jit + def _edge_env_mat_bwd_kernel( + edge_vec_ptr, # (E, 3) + ctype_ptr, # (E,) + emask_ptr, # (E,) + std_ptr, # (ntypes, 4) + genv_ptr, # (E, 4) + gsw_ptr, # (E,) upstream gradient of the switch (strip gate) + gedge_ptr, # (E, 3) output: grad wrt edge_vec + rcut, + rcut_smth, + protection, + sw_c, + n_edge, + BLOCK_E: tl.constexpr, + ): + """Closed-form ``d env / d edge_vec`` per edge; edge_vec is the leaf, so + no scatter -- each edge writes its own gradient row directly. + + The env quartet and the separately emitted switch ``sw`` share the same + radial dependence on ``|edge_vec|``, so the ``sw`` cotangent folds into + the common radial coefficient as ``g_sw * s'(L)``. The ``inv_l`` guard + (zero at a zero-length padding edge) keeps that term finite. + """ + pid = tl.program_id(0) + e = pid * BLOCK_E + tl.arange(0, BLOCK_E) + m = e < n_edge + dx = tl.load(edge_vec_ptr + e * 3 + 0, mask=m, other=0.0) + dy = tl.load(edge_vec_ptr + e * 3 + 1, mask=m, other=0.0) + dz = tl.load(edge_vec_ptr + e * 3 + 2, mask=m, other=0.0) + ct = tl.load(ctype_ptr + e, mask=m, other=0) + valid = tl.load(emask_ptr + e, mask=m, other=0) + length = tl.sqrt(dx * dx + dy * dy + dz * dz) + # safe-norm gradient: zero at a zero-length (padding) edge. + inv_l = tl.where(length > 0.0, 1.0 / length, 0.0) + length = length + tl.where(valid != 0, 0.0, 1.0) + denom = length + protection + sw, dsw = _switch(length, rcut_smth, rcut, sw_c, 0) + inv_d = 1.0 / denom + inv_d2 = inv_d * inv_d + inv_d3 = inv_d2 * inv_d + mb = ct * 4 + s0 = tl.load(std_ptr + mb + 0, mask=m, other=1.0) + s1 = tl.load(std_ptr + mb + 1, mask=m, other=1.0) + s2 = tl.load(std_ptr + mb + 2, mask=m, other=1.0) + s3 = tl.load(std_ptr + mb + 3, mask=m, other=1.0) + g0 = tl.load(genv_ptr + e * 4 + 0, mask=m, other=0.0) / s0 + gvx = tl.load(genv_ptr + e * 4 + 1, mask=m, other=0.0) / s1 + gvy = tl.load(genv_ptr + e * 4 + 2, mask=m, other=0.0) / s2 + gvz = tl.load(genv_ptr + e * 4 + 3, mask=m, other=0.0) / s3 + gsw = tl.load(gsw_ptr + e, mask=m, other=0.0) + gdot = gvx * dx + gvy * dy + gvz * dz + coef = ( + g0 * (dsw * inv_d - sw * inv_d2) + + gdot * (dsw * inv_d2 - 2.0 * sw * inv_d3) + + gsw * dsw + ) * inv_l + swq2 = sw * inv_d2 + tl.store(gedge_ptr + e * 3 + 0, coef * dx + swq2 * gvx, mask=m) + tl.store(gedge_ptr + e * 3 + 1, coef * dy + swq2 * gvy, mask=m) + tl.store(gedge_ptr + e * 3 + 2, coef * dz + swq2 * gvz, mask=m) + + +def _switch_reference( + length: Tensor, rmin: float, rmax: float, use_exp: bool +) -> tuple[Tensor, Tensor]: + """Eager smooth switch and its derivative (clamped-window semantics).""" + if use_exp: + a = _EXP_SWITCH_C / rmin + inside = (length > 0.0) & (length < rmax) + lc = length.clamp(0.0, rmax) + xarg = a * (lc - rmin) + e = torch.exp(xarg) + w = torch.exp(-e) + # -a e w = -a exp(xarg - e): the fused exponent stays finite as e -> inf, + # avoiding the 0*inf NaN of the factored form once e overflows in fp32. + dw = torch.where(inside, -a * torch.exp(xarg - e), torch.zeros_like(length)) + else: + inside = (length > rmin) & (length < rmax) + u = ((length - rmin) / (rmax - rmin)).clamp(0.0, 1.0) + u2 = u * u + w = u2 * u * (-6.0 * u2 + 15.0 * u - 10.0) + 1.0 + dw = torch.where( + inside, + (-30.0 * u2 * u2 + 60.0 * u2 * u - 30.0 * u2) / (rmax - rmin), + torch.zeros_like(length), + ) + return w, dw + + +def _geometry( + coord: Tensor, nlist: Tensor +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Shared prologue: masked neighbor index, relative vector, and safe length. + + Returns ``(mask, j_safe, d, length_safe, center_index)`` with ``d`` the + ``(nf, nloc, nnei, 3)`` relative vector and ``length_safe`` its norm with + invalid slots pinned to 1 (kept finite; nulled downstream by the mask). + """ + nf, nloc, nnei = nlist.shape + mask = nlist >= 0 + j = torch.where(mask, nlist, torch.zeros_like(nlist)) + ci = coord[:, :nloc] + cj = torch.gather(coord, 1, j.reshape(nf, -1, 1).expand(-1, -1, 3)).reshape( + nf, nloc, nnei, 3 + ) + d = cj - ci[:, :, None, :] + length = torch.linalg.norm(d, dim=-1) + length_safe = torch.where(mask, length, torch.ones_like(length)) + return mask, j, d, length_safe, ci + + +def _env_mat_reference( + coord: Tensor, + nlist: Tensor, + atype: Tensor, + mean: Tensor, + stddev: Tensor, + rcut: float, + rcut_smth: float, + radial_only: bool, + protection: float, + use_exp_switch: bool, +) -> tuple[Tensor, Tensor, Tensor]: + """Eager environment matrix, numerically identical to ``prod_env_mat``.""" + mask, _, d, length_safe, _ = _geometry(coord, nlist) + q = length_safe + protection + sw_raw, _ = _switch_reference(length_safe, rcut_smth, rcut, use_exp_switch) + w = (sw_raw * mask).unsqueeze(-1) + t0 = 1.0 / q.unsqueeze(-1) + if radial_only: + raw = t0 * w + else: + raw = torch.cat([t0, d / q.unsqueeze(-1) ** 2], dim=-1) * w + env = (raw - mean[atype]) / stddev[atype] + return env, d * mask.unsqueeze(-1), (sw_raw * mask).unsqueeze(-1) + + +def _env_mat_grad_coord_reference( + coord: Tensor, + nlist: Tensor, + atype: Tensor, + stddev: Tensor, + g_env: Tensor, + g_diff: Tensor, + g_sw: Tensor, + rcut: float, + rcut_smth: float, + radial_only: bool, + protection: float, + use_exp_switch: bool, +) -> Tensor: + """Eager closed-form ``grad_coord`` for the three env-matrix outputs.""" + nf, nloc, nnei = nlist.shape + nall = coord.shape[1] + mask, j, d, length_safe, _ = _geometry(coord, nlist) + q = length_safe + protection + sw_raw, dsw_raw = _switch_reference(length_safe, rcut_smth, rcut, use_exp_switch) + w = sw_raw * mask + wp = dsw_raw * mask + inv_l = torch.where(mask, 1.0 / length_safe, torch.zeros_like(length_safe)) + g_raw = g_env / stddev[atype] + g0 = g_raw[..., 0] + if radial_only: + gdot = torch.zeros_like(g0) + gv = torch.zeros_like(d) + else: + gv = g_raw[..., 1:4] + gdot = (gv * d).sum(-1) + coef = ( + g0 * (wp / q - w / q**2) + + gdot * (wp / q**2 - 2.0 * w / q**3) + + g_sw[..., 0] * wp + ) * inv_l + g_d = coef[..., None] * d + (w / q**2)[..., None] * gv + g_diff + g_d = g_d * mask[..., None] + + grad = torch.zeros(nf, nall, 3, dtype=coord.dtype, device=coord.device) + grad = grad.scatter_add( + 1, j.reshape(nf, -1, 1).expand(-1, -1, 3), g_d.reshape(nf, -1, 3) + ) + grad = grad.index_add( + 1, torch.arange(nloc, dtype=torch.int64, device=coord.device), -g_d.sum(2) + ) + return grad + + +def _use_triton(coord: Tensor) -> bool: + return TRITON_AVAILABLE and coord.is_cuda + + +def _launch(coord: Tensor, nlist: Tensor, nnei: int): + nf, nloc = nlist.shape[0], nlist.shape[1] + return (nf * nloc,), triton.next_power_of_2(max(nnei, 1)) + + +def _env_mat_fwd_impl( + coord: Tensor, + nlist: Tensor, + atype: Tensor, + mean: Tensor, + stddev: Tensor, + rcut: float, + rcut_smth: float, + radial_only: bool, + protection: float, + use_exp_switch: bool, +) -> tuple[Tensor, Tensor, Tensor]: + if not _use_triton(coord): + return _env_mat_reference( + coord, + nlist, + atype, + mean, + stddev, + rcut, + rcut_smth, + radial_only, + protection, + use_exp_switch, + ) + nf, nloc, nnei = nlist.shape + nall = coord.shape[1] + channels = 1 if radial_only else 4 + env = torch.empty(nf, nloc, nnei, channels, dtype=coord.dtype, device=coord.device) + diff = torch.empty(nf, nloc, nnei, 3, dtype=coord.dtype, device=coord.device) + sw = torch.empty(nf, nloc, nnei, dtype=coord.dtype, device=coord.device) + grid, block_n = _launch(coord, nlist, nnei) + wrap_triton(_env_mat_fwd_kernel)[grid]( + coord.contiguous(), + nlist.contiguous(), + atype.contiguous(), + mean.contiguous(), + stddev.contiguous(), + env, + diff, + sw, + float(rcut), + float(rcut_smth), + float(protection), + _EXP_SWITCH_C, + nloc, + nnei, + nall, + C=channels, + RADIAL=radial_only, + USE_EXP=use_exp_switch, + BLOCK_N=block_n, + ) + return env, diff, sw.unsqueeze(-1) + + +def _env_mat_grad_impl( + coord: Tensor, + nlist: Tensor, + atype: Tensor, + stddev: Tensor, + g_env: Tensor, + g_diff: Tensor, + g_sw: Tensor, + rcut: float, + rcut_smth: float, + radial_only: bool, + protection: float, + use_exp_switch: bool, +) -> Tensor: + """Coordinate gradient ``(nf, nall, 3)`` scattered from the env-mat outputs.""" + if not _use_triton(coord): + return _env_mat_grad_coord_reference( + coord, + nlist, + atype, + stddev, + g_env, + g_diff, + g_sw, + rcut, + rcut_smth, + radial_only, + protection, + use_exp_switch, + ) + nf, nloc, nnei = nlist.shape + nall = coord.shape[1] + channels = 1 if radial_only else 4 + grad = torch.zeros(nf, nall, 3, dtype=coord.dtype, device=coord.device) + grid, block_n = _launch(coord, nlist, nnei) + wrap_triton(_env_mat_bwd_kernel)[grid]( + coord.contiguous(), + nlist.contiguous(), + atype.contiguous(), + stddev.contiguous(), + g_env.contiguous(), + g_diff.contiguous(), + g_sw.contiguous(), + grad, + float(rcut), + float(rcut_smth), + float(protection), + _EXP_SWITCH_C, + nloc, + nnei, + nall, + C=channels, + RADIAL=radial_only, + USE_EXP=use_exp_switch, + BLOCK_N=block_n, + ) + return grad + + +_fwd_op = triton_op("deepmd_triton::env_mat", mutates_args=())(_env_mat_fwd_impl) +_grad_op = triton_op("deepmd_triton::env_mat_grad", mutates_args=())(_env_mat_grad_impl) + + +@_fwd_op.register_fake +def _( + coord, + nlist, + atype, + mean, + stddev, + rcut, + rcut_smth, + radial_only, + protection, + use_exp_switch, +): + nf, nloc, nnei = nlist.shape + channels = 1 if radial_only else 4 + return ( + coord.new_empty((nf, nloc, nnei, channels)), + coord.new_empty((nf, nloc, nnei, 3)), + coord.new_empty((nf, nloc, nnei, 1)), + ) + + +@_grad_op.register_fake +def _( + coord, + nlist, + atype, + stddev, + g_env, + g_diff, + g_sw, + rcut, + rcut_smth, + radial_only, + protection, + use_exp_switch, +): + return coord.new_empty((coord.shape[0], coord.shape[1], 3)) + + +def _fwd_setup_context(ctx, inputs, output): + ( + coord, + nlist, + atype, + mean, + stddev, + rcut, + rcut_smth, + radial_only, + protection, + use_exp_switch, + ) = inputs + ctx.save_for_backward(coord, nlist, atype, stddev) + ctx.cfg = (rcut, rcut_smth, radial_only, protection, use_exp_switch) + + +def _fwd_backward(ctx, g_env, g_diff, g_sw): + coord, nlist, atype, stddev = ctx.saved_tensors + rcut, rcut_smth, radial_only, protection, use_exp_switch = ctx.cfg + grad = _grad_op( + coord, + nlist, + atype, + stddev, + g_env, + g_diff, + g_sw, + rcut, + rcut_smth, + radial_only, + protection, + use_exp_switch, + ) + return grad, None, None, None, None, None, None, None, None, None + + +_fwd_op.register_autograd(_fwd_backward, setup_context=_fwd_setup_context) + + +def env_mat( + coord: Tensor, + nlist: Tensor, + atype: Tensor, + mean: Tensor, + stddev: Tensor, + rcut: float, + rcut_smth: float, + radial_only: bool = False, + protection: float = 0.0, + use_exp_switch: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + """Fused smooth environment matrix, a drop-in for ``prod_env_mat``. + + Parameters + ---------- + coord : Tensor + Extended coordinates with shape (nf, nall, 3) or (nf, nall * 3). + nlist : Tensor + Neighbor indices with shape (nf, nloc, nnei); ``< 0`` marks padding. + atype : Tensor + Local atom types with shape (nf, nloc), indexing ``mean`` / ``stddev``. + mean : Tensor + Per-type running average with shape (ntypes, nnei, 4) -- or (ntypes, + nnei, 1) when ``radial_only``. + stddev : Tensor + Per-type running standard deviation, same shape as ``mean``. + rcut : float + Neighbor cutoff radius, in the coordinate length unit. + rcut_smth : float + Inner radius where the smooth switch starts to decay. + radial_only : bool + Whether to emit only the radial channel (shape (..., 1)); otherwise the + full quartet (shape (..., 4)). + protection : float + Additive protection on the inverse distance, guarding division by zero. + use_exp_switch : bool + Whether to use the exponential switch instead of the quintic one. + + Returns + ------- + tuple of Tensor + ``(env_mat, diff, switch)`` with shapes (nf, nloc, nnei, 4 or 1), + (nf, nloc, nnei, 3), and (nf, nloc, nnei, 1). + + Notes + ----- + Routes to the Triton operator at ``DP_TRITON_INFER >= 1`` on CUDA; elsewhere + (CPU, Triton absent, or under a double-backward / training graph) it uses the + eager reference so results are identical. The registered backward is + first-order (the force path); training keeps the eager autograd chain. + + The scalar hyper-parameters are passed as kernel arguments (fp32) rather than + a device tensor: a device tensor built from Python scalars forces a + synchronizing host-to-device copy that drains the launch queue every call and + dominates the eager cost. In fp32 they are exact; in fp64 they are exact for + values representable in fp32 (typical ``rcut`` / ``rcut_smth`` / zero + ``protection``) and otherwise perturb the result at ~1e-8. + """ + if coord.dim() == 2: + coord = coord.reshape(coord.shape[0], -1, 3) + # Level gate only: the operator is called unconditionally so that it is + # captured as an opaque node under a CPU ``make_fx`` trace (the pt_expt + # freeze), while its implementation resolves the CUDA kernel vs. the eager + # reference per the runtime device. + if triton_infer_level() >= 1 and TRITON_AVAILABLE: + return _fwd_op( + coord, + nlist, + atype, + mean, + stddev, + rcut, + rcut_smth, + radial_only, + protection, + use_exp_switch, + ) + return _env_mat_reference( + coord, + nlist, + atype, + mean, + stddev, + rcut, + rcut_smth, + radial_only, + protection, + use_exp_switch, + ) + + +# ====================================================================== +# Graph-native (edge-stream) environment matrix +# ====================================================================== +# The edge form is the slot-free analogue used by the graph lower: the relative +# vector ``edge_vec = r_j - r_i`` is given directly (no neighbor gather), the +# per-type statistics are indexed by the center (dst) atom only, and the smooth +# switch is quintic. The backward is w.r.t. ``edge_vec`` (the autograd leaf on +# the graph path), so no scatter is needed -- each edge owns its gradient row. + +# Edges per program block; the kernel is memory-bound, so this is not tuned. +_EDGE_BLOCK = 256 + + +def _edge_env_mat_reference( + edge_vec: Tensor, + center_type: Tensor, + davg: Tensor, + dstd: Tensor, + rcut: float, + rcut_smth: float, + protection: float, + edge_mask: Tensor | None, +) -> tuple[Tensor, Tensor]: + """Eager per-edge environment matrix, identical to the dpmodel reference. + + Returns ``(env, sw)`` with ``sw`` (E, 1) the switch zeroed on padding + edges. The safe norm keeps the autograd force gradient finite at a + zero-length padding edge. + """ + length = safe_for_vector_norm(edge_vec, axis=-1, keepdims=True) + if edge_mask is not None: + length = length + (~edge_mask.bool())[:, None].to(length.dtype) + else: + length = torch.where(length < 1e-10, torch.ones_like(length), length) + denom = length + protection + sw, _ = _switch_reference(length, rcut_smth, rcut, False) + em = torch.cat([1.0 / denom, edge_vec / denom**2], dim=-1) * sw + env = (em - davg[center_type]) / dstd[center_type] + if edge_mask is not None: + sw = sw * edge_mask[:, None].to(sw.dtype) + return env, sw + + +def _edge_env_mat_grad_reference( + edge_vec: Tensor, + center_type: Tensor, + dstd: Tensor, + g_env: Tensor, + g_sw: Tensor, + rcut: float, + rcut_smth: float, + protection: float, + edge_mask: Tensor | None, +) -> Tensor: + """Eager closed-form ``grad_edge_vec`` for the per-edge environment matrix. + + ``g_sw`` (E, 1) is the upstream gradient of the switch output; it enters the + shared radial coefficient as ``g_sw * s'(L)`` (mirroring the fused kernel). + """ + d = edge_vec + length_raw = safe_for_vector_norm(d, axis=-1) + inv_l = torch.where( + length_raw > 0.0, 1.0 / length_raw, torch.zeros_like(length_raw) + ) + if edge_mask is not None: + length = length_raw + (~edge_mask.bool()).to(d.dtype) + else: + length = torch.where( + length_raw < 1e-10, torch.ones_like(length_raw), length_raw + ) + denom = length + protection + sw, dsw = _switch_reference(length, rcut_smth, rcut, False) + g_raw = g_env / dstd[center_type] + g0 = g_raw[:, 0] + gv = g_raw[:, 1:4] + gdot = (gv * d).sum(-1) + coef = ( + g0 * (dsw / denom - sw / denom**2) + + gdot * (dsw / denom**2 - 2.0 * sw / denom**3) + + g_sw[:, 0] * dsw + ) * inv_l + return coef[:, None] * d + (sw / denom**2)[:, None] * gv + + +def _edge_env_mat_fwd_impl( + edge_vec: Tensor, + center_type: Tensor, + edge_mask: Tensor, + davg: Tensor, + dstd: Tensor, + rcut: float, + rcut_smth: float, + protection: float, +) -> tuple[Tensor, Tensor]: + if not _use_triton(edge_vec): + return _edge_env_mat_reference( + edge_vec, center_type, davg, dstd, rcut, rcut_smth, protection, edge_mask + ) + n_edge = edge_vec.shape[0] + env = torch.empty(n_edge, 4, dtype=edge_vec.dtype, device=edge_vec.device) + sw = torch.empty(n_edge, dtype=edge_vec.dtype, device=edge_vec.device) + grid = (triton.cdiv(max(n_edge, 1), _EDGE_BLOCK),) + wrap_triton(_edge_env_mat_fwd_kernel)[grid]( + edge_vec.contiguous(), + center_type.contiguous(), + edge_mask.to(torch.int32).contiguous(), + davg.contiguous(), + dstd.contiguous(), + env, + sw, + float(rcut), + float(rcut_smth), + float(protection), + _EXP_SWITCH_C, + n_edge, + BLOCK_E=_EDGE_BLOCK, + ) + return env, sw.unsqueeze(-1) + + +def _edge_env_mat_grad_impl( + edge_vec: Tensor, + center_type: Tensor, + edge_mask: Tensor, + dstd: Tensor, + g_env: Tensor, + g_sw: Tensor, + rcut: float, + rcut_smth: float, + protection: float, +) -> Tensor: + if not _use_triton(edge_vec): + return _edge_env_mat_grad_reference( + edge_vec, + center_type, + dstd, + g_env, + g_sw, + rcut, + rcut_smth, + protection, + edge_mask, + ) + n_edge = edge_vec.shape[0] + gedge = torch.empty(n_edge, 3, dtype=edge_vec.dtype, device=edge_vec.device) + grid = (triton.cdiv(max(n_edge, 1), _EDGE_BLOCK),) + wrap_triton(_edge_env_mat_bwd_kernel)[grid]( + edge_vec.contiguous(), + center_type.contiguous(), + edge_mask.to(torch.int32).contiguous(), + dstd.contiguous(), + g_env.contiguous(), + g_sw.reshape(-1).contiguous(), + gedge, + float(rcut), + float(rcut_smth), + float(protection), + _EXP_SWITCH_C, + n_edge, + BLOCK_E=_EDGE_BLOCK, + ) + return gedge + + +_edge_fwd_op = triton_op("deepmd_triton::edge_env_mat", mutates_args=())( + _edge_env_mat_fwd_impl +) +_edge_grad_op = triton_op("deepmd_triton::edge_env_mat_grad", mutates_args=())( + _edge_env_mat_grad_impl +) + + +@_edge_fwd_op.register_fake +def _(edge_vec, center_type, edge_mask, davg, dstd, rcut, rcut_smth, protection): + return ( + edge_vec.new_empty((edge_vec.shape[0], 4)), + edge_vec.new_empty((edge_vec.shape[0], 1)), + ) + + +@_edge_grad_op.register_fake +def _(edge_vec, center_type, edge_mask, dstd, g_env, g_sw, rcut, rcut_smth, protection): + return edge_vec.new_empty((edge_vec.shape[0], 3)) + + +def _edge_setup_context(ctx, inputs, output): + edge_vec, center_type, edge_mask, davg, dstd, rcut, rcut_smth, protection = inputs + ctx.save_for_backward(edge_vec, center_type, edge_mask, dstd) + ctx.cfg = (rcut, rcut_smth, protection) + + +def _edge_backward(ctx, g_env, g_sw): + edge_vec, center_type, edge_mask, dstd = ctx.saved_tensors + rcut, rcut_smth, protection = ctx.cfg + g_edge = _edge_grad_op( + edge_vec, + center_type, + edge_mask, + dstd, + g_env, + g_sw, + rcut, + rcut_smth, + protection, + ) + return g_edge, None, None, None, None, None, None, None + + +_edge_fwd_op.register_autograd(_edge_backward, setup_context=_edge_setup_context) + + +def edge_env_mat( + edge_vec: Tensor, + center_type: Tensor, + davg: Tensor, + dstd: Tensor, + rcut: float, + rcut_smth: float, + protection: float = 0.0, + edge_mask: Tensor | None = None, + return_sw: bool = False, +) -> Tensor | tuple[Tensor, Tensor]: + """Fused per-edge environment matrix, a drop-in for the dpmodel edge form. + + Parameters + ---------- + edge_vec : Tensor + Relative vectors ``r_src - r_dst`` with shape (E, 3); padding edges must + carry a zero vector. + center_type : Tensor + Center (dst) atom type per edge with shape (E,), indexing the statistics. + davg : Tensor + Per-center-type mean with shape (ntypes, 4) (slot-independent). + dstd : Tensor + Per-center-type standard deviation, same shape as ``davg``. + rcut : float + Neighbor cutoff radius. + rcut_smth : float + Inner radius where the smooth (quintic) switch starts to decay. + protection : float + Additive protection on the inverse distance. + edge_mask : Tensor or None + Boolean valid-edge mask with shape (E,); invalid edges take + ``length + 1`` (matching the dense ``length + ~mask`` guard). When + ``None`` a zero-length threshold guard is used and the eager reference + serves the request. + return_sw : bool + When ``True``, also return the per-edge smooth switch (E, 1), zeroed on + padding edges (the strip type-pair gate consumes it). Mirrors the dense + :func:`env_mat`, which always returns the switch. + + Returns + ------- + Tensor + The per-edge environment matrix with shape (E, 4); or, when + ``return_sw`` is set, the tuple ``(env, sw)`` with ``sw`` of shape + (E, 1). + + Notes + ----- + Routes to the Triton operator at ``DP_TRITON_INFER >= 1`` when an + ``edge_mask`` is supplied (the graph path always provides one); it is called + unconditionally there so a CPU ``make_fx`` trace captures it as an opaque + node, while the implementation resolves the CUDA kernel vs. the eager + reference per the runtime device. The registered backward differentiates + ``edge_vec`` (the graph-path force leaf) and folds the switch cotangent when + ``return_sw`` feeds a downstream gradient. + """ + if triton_infer_level() >= 1 and TRITON_AVAILABLE and edge_mask is not None: + env, sw = _edge_fwd_op( + edge_vec, center_type, edge_mask, davg, dstd, rcut, rcut_smth, protection + ) + else: + env, sw = _edge_env_mat_reference( + edge_vec, center_type, davg, dstd, rcut, rcut_smth, protection, edge_mask + ) + return (env, sw) if return_sw else env diff --git a/deepmd/kernels/utils.py b/deepmd/kernels/utils.py index 08837e8fde..df2997b7d3 100644 --- a/deepmd/kernels/utils.py +++ b/deepmd/kernels/utils.py @@ -3,9 +3,9 @@ Environment-variable gates for the SeZM/DPA4 hardware-accelerated kernels. This module centralizes the opt-in selectors that route inference through the -custom Triton and CuTe kernel packages. The gates are read once at model -construction time so that they become compile-time constants in the traced -(``make_fx``) graph. +custom Triton, CuTe and hand-written CUDA kernel packages. The gates are read +once at model construction time so that they become compile-time constants in +the traced (``make_fx``) graph. """ from __future__ import ( @@ -24,27 +24,43 @@ def triton_infer_level() -> int: The level is read at module construction time so that it becomes a compile-time constant in the traced (``make_fx``) graph. It only takes - effect during inference; training always uses the dense reference path. + effect during inference. Levels are cumulative: - ``0`` -- Triton disabled; every operation uses the dense reference path. - - ``1`` -- universal kernels that need no launch-configuration table: - block-diagonal rotation, radial degree mixing, the ``SO2Linear`` - block GEMM, Wigner monomials, flash-attention aggregation, and the - segmented force assembly. These are either runtime-autotuned or run a - single shape-independent configuration. - - ``2`` -- adds kernels whose launch configuration is resolved from the - swept ``(focus_dim, lmax)`` / ``(C_wide, lmax)`` tables in - :mod:`.triton.tile_configs`: the fused SO(2) value path and the - edge-block backward kernels. A key absent from a table falls back to - the level-1 kernel (or a spill-safe configuration) for that operation, - so unswept shapes never regress below level 1. - - ``3`` -- adds the fp16x3 split-compensated mixing-stack GEMMs on - tensor cores. Entries exist only for table keys whose configuration - passed the fp64 validation sweep; unswept shapes keep the level-2 fp32 - stack. This level trades a bounded accuracy perturbation for speed - (see :mod:`.triton.so2_stack_fp16x3`). + - ``1`` -- universal kernels that need no launch-configuration table; each + either runtime-autotunes or runs a single shape-independent configuration. + + - All ``se``-family descriptors: the fused smooth environment matrix + (:mod:`.triton.env_mat`), a drop-in for the ``prod_env_mat`` front end + with a closed-form force backward. + - SeZM/DPA4: block-diagonal rotation, radial degree mixing, the + ``SO2Linear`` block GEMM, Wigner monomials, flash-attention + aggregation, and the segmented force assembly. + - DPA1 (``se_atten``): the fused environment convolution + (:mod:`.triton.dpa1.se_conv`). + + - ``2`` -- adds kernels whose launch configuration is resolved from a swept + table, falling back to the level-1 configuration so unswept shapes never + regress below level 1. + + - SeZM/DPA4: the fused SO(2) value path and the edge-block backward + kernels keyed by ``(focus_dim, lmax)`` / ``(C_wide, lmax)`` in + :mod:`.triton.sezm.tile_configs`. + - DPA1: the environment convolution keyed by ``(ng, H1)`` in + :mod:`.triton.dpa1.tile_configs`. + + - ``3`` -- adds fp16x3 split-compensated GEMMs, which recover near-fp32 + accuracy on tensor cores and trade a bounded, validated accuracy + perturbation for speed. Entries exist only for table keys whose + configuration passed the fp64 validation sweep; unswept shapes keep the + level-2 fp32 path. + + - SeZM/DPA4: the mixing-stack GEMMs + (:mod:`.triton.sezm.so2_stack_fp16x3`). + - DPA1 (``se_atten``): the compute-bound embedding last-layer GEMM + (:mod:`.triton.dpa1.gemm_fp16x3`). Returns ------- @@ -70,6 +86,65 @@ def triton_infer_level() -> int: return level +CUDA_INFER_LEVELS = (0, 1, 2) + + +def cuda_infer_level() -> int: + """Return the opt-in CUDA mega-kernel inference level from ``DP_CUDA_INFER``. + + Read at model construction time so that it becomes a compile-time constant in + the traced (``make_fx``) graph, independent of ``DP_TRITON_INFER``. It only + takes effect during inference. + + Levels are cumulative: + + - ``0`` -- CUDA mega kernels disabled; every operation uses the Triton or + dense reference path. + - ``1`` -- the fused graph-lower operator suite (separate operators, force + from ``autograd.grad``): + + - DPA1 (``se_atten``): the descriptor mega kernels + (:mod:`.cuda.dpa1.graph_descriptor`) serve the concat, + attention-free graph lower, and the energy fitting runs through the + fused cuBLAS network (:mod:`.cuda.graph_fitting`). + - All graph-lowered models: the force / virial assembly scatters + through :mod:`.cuda.edge_force_virial`. + + - ``2`` -- adds the end-to-end energy-force operator: a graph-lowered energy + model whose descriptor and fitting are both fused-eligible collapses its + descriptor, fitting and analytic force / virial assembly into one operator + that returns the force as a value (no autograd tape), numerically + identical to level 1. A model outside that class falls back to the level-1 + operators, so level 2 never regresses below level 1. + + - DPA1 (``se_atten``): the attention-free graph lower with a + fused-eligible energy fitting routes through + :mod:`.cuda.dpa1.graph_energy_force`. + + Returns + ------- + int + The configured level in ``{0, 1, 2}``. + + Raises + ------ + ValueError + If ``DP_CUDA_INFER`` is not an integer in ``{0, 1, 2}``. + """ + raw = os.environ.get("DP_CUDA_INFER", "0").strip() + try: + level = int(raw) + except ValueError: + raise ValueError( + f"DP_CUDA_INFER must be an integer in {CUDA_INFER_LEVELS}, got {raw!r}" + ) from None + if level not in CUDA_INFER_LEVELS: + raise ValueError( + f"DP_CUDA_INFER must be one of {CUDA_INFER_LEVELS}, got {level}" + ) + return level + + def use_cute_infer() -> bool: """Return whether the opt-in CuTe inference operator is enabled. diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py index 0ffdbb7dbb..c0485b5746 100644 --- a/deepmd/pt/model/descriptor/env_mat.py +++ b/deepmd/pt/model/descriptor/env_mat.py @@ -2,6 +2,13 @@ import torch +from deepmd.kernels.triton.env_mat import ( + TRITON_AVAILABLE, +) +from deepmd.kernels.triton.env_mat import env_mat as _env_mat_triton +from deepmd.kernels.utils import ( + triton_infer_level, +) from deepmd.pt.utils.preprocess import ( compute_exp_sw, compute_smooth_weight, @@ -77,6 +84,26 @@ def prod_env_mat( ------- - env_mat: Shape is [nframes, natoms[1]*nnei*4]. """ + # Opt-in inference (``DP_TRITON_INFER >= 1``, CUDA): the fused Triton kernel + # forms the environment matrix in one node-parallel pass and carries a + # closed-form backward for the force path. Training (level 0) and the CPU + # path keep the dense autograd chain below, which supports higher-order + # differentiation. The block is nested under ``torch.jit.is_scripting`` so + # the whole (non-scriptable) branch is pruned under ``torch.jit.script``. + if not torch.jit.is_scripting(): + if TRITON_AVAILABLE and triton_infer_level() >= 1 and extended_coord.is_cuda: + return _env_mat_triton( + extended_coord, + nlist, + atype, + mean, + stddev, + rcut, + rcut_smth, + radial_only=radial_only, + protection=protection, + use_exp_switch=use_exp_switch, + ) _env_mat_se_a, diff, switch = _make_env_mat( nlist, extended_coord, diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 32c8830b1e..808515702a 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -13,6 +13,15 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.kernels.triton.dpa1.activation import ( + ACT_CODES, +) +from deepmd.kernels.triton.dpa1.se_conv import ( + se_atten_conv, +) +from deepmd.kernels.utils import ( + triton_infer_level, +) from deepmd.pt.model.descriptor.descriptor import ( DescriptorBlock, ) @@ -290,6 +299,24 @@ def __init__( "type_embd_data", torch.zeros(0, dtype=self.prec, device=env.DEVICE) ) + # Opt-in fused environment-convolution routing (inference only, gated by + # ``DP_TRITON_INFER``). The fused kernel folds the last embedding layer + # (its ``tanh`` activation, timestep and residual, doubling or identity), + # the type-pair gate (strip only), the smooth cutoff and the moment + # reduction into one node-parallel pass; the two embedding GEMMs stay on + # cuBLAS. It requires no attention layer and a last-layer activation the + # kernel inlines (``tanh`` or ``silu``, forward and backward); arbitrary + # channel widths are handled by masked padding inside the kernel. Both + # ``strip`` (type-pair gate) and ``concat`` (type feature folded into the + # embedding input, no gate) are supported. Any other configuration keeps + # the dense reference path. + self.use_triton_infer = triton_infer_level() >= 1 + self.se_conv_eligible = ( + self.tebd_input_mode in ("strip", "concat") + and self.attn_layer == 0 + and self.filter_layers.networks[0].layers[-1].activate_name in ACT_CODES + ) + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.rcut @@ -592,6 +619,10 @@ def forward( rr = dmatrix rr = rr * exclude_mask[:, :, None] ss = rr[:, :, :1] + # Whether the pair representation ``g2`` is produced this pass. The + # compressed and fused strip paths form only the moment tensor and + # leave ``g2`` empty; the flag drives the return below. + g2_is_none = self.geo_compress if self.tebd_input_mode in ["concat"]: atype_tebd_ext = extended_atype_embd # nb x (nloc x nnei) x nt @@ -613,16 +644,45 @@ def forward( else: # nfnl x nnei x (1 + tebd_dim) ss = torch.concat([ss, nlist_tebd], dim=2) - # nfnl x nnei x ng - gg = self.filter_layers.networks[0](ss) - input_r = torch.nn.functional.normalize( - rr.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1 - ) - gg = self.dpa1_attention( - gg, nlist_mask, input_r=input_r, sw=sw - ) # shape is [nframes*nloc, self.neei, out_size] - # nfnl x 4 x ng - xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) + if ( + not torch.jit.is_scripting() + and self.use_triton_infer + and self.se_conv_eligible + and not self.training + and not self.geo_compress + and rr.is_cuda + ): + # Fused environment convolution (concat, opt-in inference path). + # The type feature already enters through ``ss`` (the concat + # embedding input), so no type gate is applied; the final + # embedding layer (its timestep and residual) and the moment + # reduction collapse into one node-parallel Triton kernel while + # the two embedding GEMMs stay on cuBLAS. + xyz_scatter = se_atten_conv( + self.filter_layers.networks[0], + ss, + None, + None, + None, + rr, + gated=0, + ) + # ``gg`` (the pair representation g2) is not formed on this path; + # it is returned as ``None``, so a zero-element placeholder keeps + # ``gg`` a tensor without the full (nf, nloc, nnei, ng) allocation. + gg = xyz_scatter.new_empty(0) + g2_is_none = True + else: + # nfnl x nnei x ng + gg = self.filter_layers.networks[0](ss) + input_r = torch.nn.functional.normalize( + rr.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1 + ) + gg = self.dpa1_attention( + gg, nlist_mask, input_r=input_r, sw=sw + ) # shape is [nframes*nloc, self.neei, out_size] + # nfnl x 4 x ng + xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) elif self.tebd_input_mode in ["strip"]: assert self.filter_layers_strip is not None assert type_embedding is not None @@ -632,82 +692,108 @@ def forward( nlist_index = nlist.reshape(nb, nloc * nnei) # nf x (nl x nnei) nei_type = torch.gather(extended_atype, dim=1, index=nlist_index) + # Per-edge row index into the (padded) type-pair embedding table. if self.type_one_side: if self.tebd_compress: tt_full = self.type_embd_data else: # (ntypes+1, tebd_dim) -> (ntypes+1, ng) tt_full = self.filter_layers_strip.networks[0](type_embedding) - # (nf*nl*nnei,) -> (nf*nl*nnei, ng) - gg_t = tt_full[nei_type.view(-1).type(torch.long)] + tebd_idx = nei_type.view(-1).to(torch.long) else: idx_i = torch.tile( atype.reshape(-1, 1) * ntypes_with_padding, [1, nnei] ).view(-1) - idx_j = nei_type.view(-1) - # (nf x nl x nnei) - idx = (idx_i + idx_j).to(torch.long) + tebd_idx = (idx_i + nei_type.view(-1)).to(torch.long) if self.tebd_compress: # ((ntypes+1)^2, ng) tt_full = self.type_embd_data else: - # ((ntypes+1)^2) * (ntypes+1)^2 * nt type_embedding_nei = torch.tile( type_embedding.view(1, ntypes_with_padding, nt), [ntypes_with_padding, 1, 1], ) - # (ntypes+1)^2 * ((ntypes+1)^2) * nt type_embedding_center = torch.tile( type_embedding.view(ntypes_with_padding, 1, nt), [1, ntypes_with_padding, 1], ) - # ((ntypes+1)^2 * (ntypes+1)^2) * (nt+nt) two_side_type_embedding = torch.cat( [type_embedding_nei, type_embedding_center], -1 ).reshape(-1, nt * 2) tt_full = self.filter_layers_strip.networks[0]( two_side_type_embedding ) - # (nf x nl x nnei) x ng - gg_t = tt_full[idx] - # (nf x nl) x nnei x ng - gg_t = gg_t.reshape(nfnl, nnei, ng) - if self.smooth: - gg_t = gg_t * sw.reshape(-1, self.nnei, 1) - if self.geo_compress: - ss = ss.reshape(-1, 1) - gg_t = gg_t.reshape(-1, gg_t.size(-1)) - xyz_scatter = torch.ops.deepmd.tabulate_fusion_se_atten( - self.compress_data[0].contiguous(), - self.compress_info[0].cpu().contiguous(), - ss.contiguous(), - rr.contiguous(), - gg_t.contiguous(), - self.filter_neuron[-1], - self.is_sorted, - )[0] - # to make torchscript happy - gg = torch.empty( - nframes, - nloc, - self.nnei, - self.filter_neuron[-1], - dtype=gg_t.dtype, - device=gg_t.device, + if ( + not torch.jit.is_scripting() + and self.use_triton_infer + and self.se_conv_eligible + and not self.training + and not self.geo_compress + and rr.is_cuda + ): + # Fused environment convolution (opt-in inference path). The + # final embedding layer (its timestep and residual), the + # type-pair gate, the smooth cutoff and the moment reduction + # collapse into one node-parallel Triton kernel; the two + # embedding GEMMs remain on cuBLAS. The + # ``torch.jit.is_scripting`` guard keeps the + # TorchScript freeze (LAMMPS) on the dense reference path, since + # the operator is a ``triton_op`` composable only under eager + # and ``make_fx`` / ``torch.compile``. + sw_eff = sw if self.smooth else torch.ones_like(sw) + xyz_scatter = se_atten_conv( + self.filter_layers.networks[0], + ss, + tt_full, + tebd_idx, + sw_eff.reshape(nfnl, self.nnei), + rr, + gated=1, ) + # ``gg`` (the pair representation g2) is not formed on this path; + # it is returned as ``None``, so a zero-element placeholder keeps + # ``gg`` a tensor without the full (nf, nloc, nnei, ng) allocation. + gg = xyz_scatter.new_empty(0) + g2_is_none = True else: - # nfnl x nnei x ng - gg_s = self.filter_layers.networks[0](ss) - # nfnl x nnei x ng - gg = gg_s * gg_t + gg_s - input_r = torch.nn.functional.normalize( - rr.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1 - ) - gg = self.dpa1_attention( - gg, nlist_mask, input_r=input_r, sw=sw - ) # shape is [nframes*nloc, self.neei, out_size] - # nfnl x 4 x ng - xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) + # (nf x nl) x nnei x ng + gg_t = tt_full[tebd_idx].reshape(nfnl, nnei, ng) + if self.smooth: + gg_t = gg_t * sw.reshape(-1, self.nnei, 1) + if self.geo_compress: + ss = ss.reshape(-1, 1) + gg_t = gg_t.reshape(-1, gg_t.size(-1)) + xyz_scatter = torch.ops.deepmd.tabulate_fusion_se_atten( + self.compress_data[0].contiguous(), + self.compress_info[0].cpu().contiguous(), + ss.contiguous(), + rr.contiguous(), + gg_t.contiguous(), + self.filter_neuron[-1], + self.is_sorted, + )[0] + # to make torchscript happy + gg = torch.empty( + nframes, + nloc, + self.nnei, + self.filter_neuron[-1], + dtype=gg_t.dtype, + device=gg_t.device, + ) + else: + # nfnl x nnei x ng + gg_s = self.filter_layers.networks[0](ss) + # nfnl x nnei x ng + gg = gg_s * gg_t + gg_s + input_r = torch.nn.functional.normalize( + rr.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1 + ) + gg = self.dpa1_attention( + gg, nlist_mask, input_r=input_r, sw=sw + ) # shape is [nframes*nloc, self.neei, out_size] + # nfnl x 4 x ng + xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) else: raise NotImplementedError @@ -722,7 +808,7 @@ def forward( return ( result.view(nframes, nloc, self.filter_neuron[-1] * self.axis_neuron), gg.view(nframes, nloc, self.nnei, self.filter_neuron[-1]) - if not self.geo_compress + if not g2_is_none else None, dmatrix.view(nframes, nloc, self.nnei, 4)[..., 1:], rot_mat.view(nframes, nloc, self.filter_neuron[-1], 3), diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index a0ea96b2eb..c1520402fb 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import dataclasses from typing import ( Any, ) @@ -12,6 +13,41 @@ from deepmd.dpmodel.utils.env_mat_stat import ( merge_env_stat, ) +from deepmd.kernels.cuda.dpa1.graph_compress import ( + dpa1_graph_compress, +) +from deepmd.kernels.cuda.dpa1.graph_compress import ( + op_available as cuda_compress_available, +) +from deepmd.kernels.cuda.dpa1.graph_descriptor import ( + dpa1_graph_descriptor, +) +from deepmd.kernels.cuda.dpa1.graph_descriptor import ( + op_available as cuda_descriptor_available, +) +from deepmd.kernels.triton.dpa1.activation import ( + ACT_CODES, + TRITON_AVAILABLE, +) +from deepmd.kernels.triton.dpa1.edge_conv import ( + concat_gate_placeholders as edge_concat_gate_placeholders, +) +from deepmd.kernels.triton.dpa1.edge_conv import ( + edge_conv, +) +from deepmd.kernels.triton.dpa1.gemm_fp16x3 import ( + embed_last_gemm, +) +from deepmd.kernels.triton.dpa1.se_conv import ( + concat_gate_placeholders, + se_conv, +) +from deepmd.kernels.triton.env_mat import edge_env_mat as _edge_env_mat_triton +from deepmd.kernels.triton.env_mat import env_mat as _env_mat_triton +from deepmd.kernels.utils import ( + cuda_infer_level, + triton_infer_level, +) from deepmd.pt_expt.common import ( torch_module, ) @@ -23,6 +59,214 @@ ) +def _has_graph_fields(graph: Any, fields: tuple[str, ...]) -> bool: + return all(getattr(graph, field, None) is not None for field in fields) + + +_DESTINATION_CSR_FIELDS = ("destination_order", "destination_row_ptr") +_DUAL_CSR_FIELDS = ( + *_DESTINATION_CSR_FIELDS, + "source_order", + "source_row_ptr", +) + + +# ====================================================================== +# Fused Triton environment convolution (attn-free inference path) +# +# The prologue / epilogue below are module-level free functions taking the +# descriptor as ``desc``: pure computation shared by the tabulated +# (:meth:`DescrptDPA1._call_compressed`) and Triton +# (:meth:`DescrptDPA1._call_triton`) entry methods, across both tebd-input +# modes. ``_env_mat`` is the shared environment-matrix prologue; strip adds the +# type-pair gate through ``_strip_pair_index`` while concat folds the type +# feature into the embedding input through ``_concat_embedding_input``. The +# entry methods and the eligibility test live on the descriptor (and are +# delegated by ``DescrptSeAttenV2``), keeping the acceleration paths symmetric. +# Routing is device-free (decided from ``triton_infer_level`` and the layer +# configuration, never a tensor's device), so a CPU ``make_fx`` trace bakes the +# ``se_conv`` operator into the exported graph for the pt_expt ``.pt2``. +# ====================================================================== +def _env_mat( + desc: Any, + coord_ext: torch.Tensor, + atype_ext: torch.Tensor, + nlist: torch.Tensor, +) -> tuple: + """Environment-matrix prologue shared by every fused path (strip / concat). + + Returns ``(nf, nloc, nnei, ng, nfnl, rr, ss, sw, nlist_masked, + type_embedding)``: ``rr`` the ``(nfnl, nnei, 4)`` environment matrix + (excluded edges zeroed), ``ss`` its radial channel ``(nfnl, nnei, 1)``, + ``sw`` the ``(nfnl, nnei, 1)`` smooth cutoff (zeroed on excluded/padding + edges), and ``nlist_masked`` the neighbor indices with excluded/padding + entries mapped to ``0`` (for downstream gathers). + """ + se = desc.se_atten + nf, nloc, nnei = nlist.shape + if triton_infer_level() >= 1: + # Fused env-matrix operator, captured opaquely under the pt_expt trace and + # resolving to the Triton kernel at CUDA runtime; identical outputs to the + # array-API ``EnvMat.call`` below. + rr, _diff, sw = _env_mat_triton( + coord_ext, + nlist, + atype_ext[:, :nloc], + se.mean[...], + se.stddev[...], + se.env_mat.rcut, + se.env_mat.rcut_smth, + radial_only=False, + protection=se.env_mat.protection, + use_exp_switch=se.env_mat.use_exp_switch, + ) + else: + rr, _diff, sw = se.env_mat.call( + coord_ext, atype_ext, nlist, se.mean[...], se.stddev[...] + ) + nf, nloc, nnei, _ = rr.shape + ng = se.neuron[-1] + nfnl = nf * nloc + + exclude_mask = ( + se.emask.build_type_exclude_mask(nlist, atype_ext) + .view(nfnl, nnei) + .to(torch.bool) + ) + nlist = nlist.view(nfnl, nnei) + nlist = torch.where(exclude_mask, nlist, torch.full_like(nlist, -1)) + nlist_mask = nlist != -1 + nlist_masked = torch.where(nlist_mask, nlist, torch.zeros_like(nlist)) + sw = torch.where( + nlist_mask[:, :, None], + sw.view(nfnl, nnei, 1), + torch.zeros(nfnl, nnei, 1, dtype=sw.dtype, device=sw.device), + ) + rr = rr.view(nfnl, nnei, 4) * exclude_mask[:, :, None].to(rr.dtype) + ss = rr[:, :, :1] + + type_embedding = desc.type_embedding.call() + return nf, nloc, nnei, ng, nfnl, rr, ss, sw, nlist_masked, type_embedding + + +def _strip_pair_index( + desc: Any, + atype_ext: torch.Tensor, + nlist_masked: torch.Tensor, + type_embedding: torch.Tensor, + nf: int, + nloc: int, + nnei: int, +) -> torch.Tensor: + """Strip-mode per-edge row index into the type-pair embedding table. + + One-side uses the neighbor type; two-side folds the ``(center, neighbor)`` + pair into a flat index, mirroring the dense reference. Shape ``(nfnl*nnei,)``. + """ + se = desc.se_atten + ntypes_with_padding = type_embedding.shape[0] + nlist_index = nlist_masked.view(nf, nloc * nnei) + nei_type = torch.gather(atype_ext, dim=1, index=nlist_index) + if se.type_one_side: + return nei_type.reshape(-1).to(torch.long) + atype = atype_ext[:, :nloc] + idx_i = torch.tile(atype.reshape(-1, 1) * ntypes_with_padding, [1, nnei]).view(-1) + return (idx_i + nei_type.reshape(-1)).to(torch.long) + + +def _concat_embedding_input( + desc: Any, + ss: torch.Tensor, + atype_ext: torch.Tensor, + nlist_masked: torch.Tensor, + type_embedding: torch.Tensor, + nf: int, + nloc: int, + nnei: int, +) -> torch.Tensor: + """Concat-mode embedding input: radial channel concatenated with the + neighbor (and, two-side, center) type embeddings. + + Returns ``(nfnl, nnei, 1 + k * tebd_dim)`` with ``k in {1, 2}`` (one-side / + two-side), matching the dense concat reference exactly. + """ + se = desc.se_atten + tebd_dim = desc.tebd_dim + nfnl = nf * nloc + atype_embd_ext = type_embedding[atype_ext.reshape(-1)].view( + nf, atype_ext.shape[1], tebd_dim + ) + # neighbor type embedding, gathered per edge + nlist_index = nlist_masked.view(nf, nloc * nnei) + idx = nlist_index[:, :, None].expand(-1, -1, tebd_dim) + nlist_tebd = torch.gather(atype_embd_ext, dim=1, index=idx).view( + nfnl, nnei, tebd_dim + ) + if se.type_one_side: + return torch.cat([ss, nlist_tebd], dim=-1) + # center type embedding, broadcast over neighbors + atype_embd = atype_embd_ext[:, :nloc, :].reshape(nfnl, 1, tebd_dim) + atype_tebd = atype_embd.expand(-1, nnei, -1) + return torch.cat([ss, nlist_tebd, atype_tebd], dim=-1) + + +def _grrg_from_moment( + desc: Any, + xyz_scatter: torch.Tensor, + type_embedding: torch.Tensor, + coord_ext: torch.Tensor, + atype_ext: torch.Tensor, + nf: int, + nloc: int, + nnei: int, + ng: int, + sw: torch.Tensor, +) -> Any: + """Strip-mode epilogue: symmetry-invariant contraction of the moment. + + Consumes the unnormalized moment ``xyz_scatter`` (nfnl, 4, ng), applies the + ``1 / nnei`` normalization, forms the ``G^T G`` descriptor and the rotation + matrix, and appends the center type embedding when ``concat_output_tebd``. + """ + se = desc.se_atten + xyz_scatter = xyz_scatter / se.nnei + xyz_scatter_1 = xyz_scatter.permute(0, 2, 1) + rot_mat = xyz_scatter_1[:, :, 1:4] + xyz_scatter_2 = xyz_scatter[:, :, 0 : se.axis_neuron] + result = torch.matmul(xyz_scatter_1, xyz_scatter_2) + result = result.view(nf, nloc, ng * se.axis_neuron) + rot_mat = rot_mat.view(nf, nloc, ng, 3) + if desc.concat_output_tebd: + nall = coord_ext.view(nf, -1).shape[1] // 3 + atype_embd_ext = type_embedding[atype_ext.reshape(-1)].view( + nf, nall, desc.tebd_dim + ) + atype_embd = atype_embd_ext[:, :nloc, :] + result = torch.cat([result, atype_embd.view(nf, nloc, desc.tebd_dim)], dim=-1) + return result, rot_mat, None, None, sw.view(nf, nloc, nnei, 1) + + +def _type_pair_table(desc: Any, type_embedding: torch.Tensor) -> torch.Tensor: + """Type-pair embedding table (P, ng) gathered per edge by ``tebd_idx``. + + One-side keeps the neighbor-type rows; two-side forms every + ``(neighbor, center)`` pair, mirroring the dense reference. + """ + se = desc.se_atten + if se.type_one_side: + return se.cal_g_strip(type_embedding, 0) + nt = type_embedding.shape[1] + ntypes_with_padding = type_embedding.shape[0] + nei = type_embedding.view(1, ntypes_with_padding, nt).expand( + ntypes_with_padding, ntypes_with_padding, nt + ) + center = type_embedding.view(ntypes_with_padding, 1, nt).expand( + ntypes_with_padding, ntypes_with_padding, nt + ) + two_side = torch.cat([nei, center], dim=-1).reshape(-1, nt * 2) + return se.cal_g_strip(two_side, 0) + + @BaseDescriptor.register("se_atten") @BaseDescriptor.register("dpa1") @torch_module @@ -186,128 +430,663 @@ def call( comm_dict: dict | None = None, charge_spin: torch.Tensor | None = None, ) -> Any: - if not self.compress: - return DescrptDPA1DP.call.__wrapped__( - self, coord_ext, atype_ext, nlist, mapping + # Compressed descriptors take the tabulated dense path. + if self.compress: + return self._call_compressed(coord_ext, atype_ext, nlist) + # Fused Triton convolution (strip / concat) for inference when eligible. + if ( + triton_infer_level() >= 1 + and not self.training + and self._fused_eligible("triton") + ): + return self._call_triton(coord_ext, atype_ext, nlist) + # Array-API reference (dpmodel). + return DescrptDPA1DP.call.__wrapped__( + self, coord_ext, atype_ext, nlist, mapping + ) + + def call_graph( + self, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor | None = None, + static_nnei: int | None = None, + ) -> Any: + """Graph-native forward, routed through the fused edge kernel when eligible. + + Concat and strip attention-free blocks are all graph-eligible. A + geometrically compressed descriptor takes the tabulated route + (:meth:`_call_graph_cuda_compress` at ``DP_CUDA_INFER >= 1``, the + original fused table operator through + :meth:`_call_graph_compress_reference` at level 0); an uncompressed + descriptor takes the embedding route (:meth:`_call_graph_cuda` mega + kernel, the Triton edge convolution for concat, or the dpmodel + reference). ``static_nnei`` is the shape-static neighbor count the dense + :meth:`call` adapter supplies for the attention edge-pair enumeration; + it is forwarded to the dpmodel reference and unused by the + attention-free fused paths. + + Unlike the dense :meth:`call`, the graph lower is not + ``@cast_precision``-wrapped: ``edge_vec`` arrives in the model-agnostic + fp64 ``.pt2`` ABI. The fused CUDA kernels consume it in that leaf dtype + directly (casting to the model precision in-kernel); the Triton and + reference paths need the Python-side alignment below so their + environment matrix runs at the model precision. Force / virial precision + is decided downstream by the model output transform, not here. + """ + # Fused kernels apply to inference with a type embedding; backend + # capability is decided by :meth:`_fused_eligible`, and this method only + # routes among the eligible implementations. + fused = not self.training and type_embedding is not None + # CUDA consumes the fp64 leaf directly. Eligibility already encodes the + # compressed / MLP rule split; dispatch selects the matching kernel. + cuda_graph_eligible = not self.geo_compress or _has_graph_fields( + graph, _DESTINATION_CSR_FIELDS + ) + if ( + fused + and cuda_graph_eligible + and cuda_infer_level() >= 1 + and self._fused_eligible("cuda") + ): + if self.geo_compress: + return self._call_graph_cuda_compress(graph, atype, type_embedding) + return self._call_graph_cuda(graph, atype, type_embedding) + # Every remaining path evaluates the environment matrix in the model + # (statistics) precision, so align the fp64 leaf once. + prec = self.se_atten.mean.dtype + if graph.edge_vec.dtype != prec: + graph = dataclasses.replace(graph, edge_vec=graph.edge_vec.to(prec)) + # Triton edge convolution: serves the uncompressed concat / strip lower; + # a tabulated (geo_compress) descriptor stays on the table path below. + if ( + fused + and triton_infer_level() >= 1 + and not self.geo_compress + and self._fused_eligible("triton") + ): + return self._call_graph_triton(graph, atype, type_embedding) + # Reference: the tabulated table operator for a compressed descriptor + # (level 0 / trace-time / CPU fallback), else the dpmodel array API. + if self.geo_compress: + return self._call_graph_compress_reference(graph, atype, type_embedding) + return DescrptDPA1DP.call_graph( + self, graph, atype, type_embedding=type_embedding, static_nnei=static_nnei + ) + + def _fused_eligible(self, backend: str) -> bool: + """Whether a fused descriptor kernel can serve this block. + + Parameters + ---------- + backend : str + ``"cuda"`` for the mega kernel (:mod:`.cuda.dpa1.graph_descriptor`) + or ``"triton"`` for the environment convolution + (:mod:`.triton.dpa1`). Both require the attention-free path, both + tebd input modes (``concat`` / ``strip``) and an activation the + kernel inlines (``tanh`` / ``silu``; see + :data:`.activation.ACT_CODES`). ``cuda`` additionally requires no + excluded types, fp32 statistics, the compiled operator library, + and -- since its MLP runs entirely in registers with + template-specialized widths -- a three-layer embedding where each + layer keeps or doubles the previous width, with + ``neuron[0] in {8, 16, 32, 64}``, ``neuron[1] <= 64`` and + ``neuron[2] <= 128``, and one inlined activation on every layer + (the strip gate network is unconstrained: its table is + precomputed outside the kernel). ``triton`` serves any layer + stack (the head layers run on cuBLAS), so only the last + activation must be inlined. + """ + se = self.se_atten + if se.attn_layer != 0: + return False + layers = se.embeddings[0].layers + if backend == "cuda": + if self.geo_compress: + # The tabulated kernel replaces the embedding MLP, so its width + # stack is unconstrained; only the table width (``neuron[-1]``) + # matters, served up to 256 by the padded bucket dispatch. + return ( + se.tebd_input_mode == "strip" + and not se.exclude_types + and se.mean.dtype == torch.float32 + and self.compress_data[0].dtype == torch.float32 + and self.type_embd_data.dtype == torch.float32 + and int(se.neuron[-1]) <= 256 + and 0 < int(se.axis_neuron) <= min(16, int(se.neuron[-1])) + and cuda_compress_available() + ) + widths = [int(layer.w.shape[1]) for layer in layers] + first = layers[0] + first_has_residual = bool(first.resnet) and first.w.shape[1] in ( + first.w.shape[0], + 2 * first.w.shape[0], + ) + parameters = [ + tensor + for layer in layers + for tensor in (layer.w, layer.b, layer.idt) + if tensor is not None + ] + return ( + not self.compress + and se.tebd_input_mode in ("concat", "strip") + and not se.exclude_types + and se.mean.dtype == torch.float32 + and all(tensor.dtype == torch.float32 for tensor in parameters) + and not first_has_residual + and len(widths) == 3 + and widths[0] in (8, 16, 32, 64) + and widths[1] in (widths[0], 2 * widths[0]) + and widths[2] in (widths[1], 2 * widths[1]) + and widths[1] <= 64 + and widths[2] <= 128 + and len({str(layer.activation_function).lower() for layer in layers}) + == 1 + and str(layers[0].activation_function).lower() in ACT_CODES + and cuda_descriptor_available() ) - return self._call_compressed(coord_ext, atype_ext, nlist) + return ( + TRITON_AVAILABLE + and se.tebd_input_mode in ("strip", "concat") + and str(layers[-1].activation_function).lower() in ACT_CODES + ) - def _call_compressed( + def _call_triton( self, coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, ) -> Any: - """Compressed forward for DPA1 descriptor.""" - # env_mat: nf x nloc x nnei x 4 - rr, _diff, sw = self.se_atten.env_mat.call( + """Fused Triton environment convolution (attn-free strip / concat path). + + The two embedding GEMMs stay on cuBLAS; the last layer's activation and + residual, the smooth cutoff and the moment reduction collapse into + :func:`se_conv`, which never materializes an ``(E, ng)`` tensor. Strip + additionally folds the type-pair gate (gathered inline); concat instead + feeds the type feature through the embedding input and applies no gate. + Composes under ``make_fx`` / ``torch.export`` so the operator is baked + into the pt_expt ``.pt2``. + """ + nf, nloc, nnei, ng, nfnl, rr, ss, sw, nlist_masked, type_embedding = _env_mat( + self, coord_ext, atype_ext, nlist + ) + se = self.se_atten + strip = se.tebd_input_mode == "strip" + # Embedding-net input: the radial channel (strip) or the radial-plus- + # type-embedding concatenation (concat). + if strip: + emb_in = ss + else: + emb_in = _concat_embedding_input( + self, ss, atype_ext, nlist_masked, type_embedding, nf, nloc, nnei + ) + + # Evaluate the embedding net through its penultimate layer (cuBLAS) and + # fold the final layer's pre-activation into se_conv. + *head, last = se.embeddings[0].layers + h = emb_in + for layer in head: + h = layer.call(h) + z2 = embed_last_gemm(h, last.w, last.b) + h1_dim, out_dim = last.w.shape + if last.resnet and out_dim == 2 * h1_dim: + resnet_mult = 2 + elif last.resnet and out_dim == h1_dim: + resnet_mult = 1 + else: + resnet_mult = 0 + idt = ( + last.idt + if last.idt is not None + else torch.ones(out_dim, dtype=z2.dtype, device=z2.device) + ) + act = ACT_CODES[str(last.activation_function).lower()] + if strip: + gated = 1 + tt_full = _type_pair_table(self, type_embedding) + tebd_idx = _strip_pair_index( + self, atype_ext, nlist_masked, type_embedding, nf, nloc, nnei + ) + sw_eff = ( + sw.view(nfnl, nnei) + if se.smooth + else torch.ones(nfnl, nnei, dtype=rr.dtype, device=rr.device) + ) + else: + # Concat folds the type feature into ``emb_in``; the gate inputs are + # unused by the kernel (``gated == 0``). + gated = 0 + tt_full, tebd_idx, sw_eff = concat_gate_placeholders(z2, ng) + # Unnormalized moment (nfnl, 4, ng); _grrg_from_moment applies 1 / nnei. + xyz_scatter = se_conv( + z2.contiguous(), + h.contiguous(), + idt, + tt_full, + tebd_idx, + sw_eff, + rr, + resnet_mult, + act, + gated, + ) + return _grrg_from_moment( + self, + xyz_scatter, + type_embedding, coord_ext, atype_ext, - nlist, - self.se_atten.mean[...], - self.se_atten.stddev[...], + nf, + nloc, + nnei, + ng, + sw, ) - nf, nloc, nnei, _ = rr.shape - ng = self.se_atten.neuron[-1] - nfnl = nf * nloc - - # Exclude mask and nlist processing - exclude_mask = self.se_atten.emask.build_type_exclude_mask(nlist, atype_ext) - exclude_mask = exclude_mask.view(nfnl, nnei) - nlist = nlist.view(nfnl, nnei) - exclude_mask = exclude_mask.to(torch.bool) - nlist = torch.where(exclude_mask, nlist, torch.full_like(nlist, -1)) - nlist_mask = nlist != -1 - nlist_masked = torch.where(nlist_mask, nlist, torch.zeros_like(nlist)) - # nfnl x nnei x 1 - sw = torch.where( - nlist_mask[:, :, None], - sw.view(nfnl, nnei, 1), - torch.zeros(nfnl, nnei, 1, dtype=sw.dtype, device=sw.device), - ) - - # nfnl x nnei x 4 - rr = rr.view(nfnl, nnei, 4) - rr = rr * exclude_mask[:, :, None].to(rr.dtype) - # nfnl x nnei x 1 - ss = rr[:, :, :1] - # Type embedding lookup from precomputed buffer - type_embedding = self.type_embedding.call() - ntypes_with_padding = type_embedding.shape[0] - # nf x (nloc x nnei) - nlist_index = nlist_masked.view(nf, nloc * nnei) - # nf x (nloc x nnei) - nei_type = torch.gather(atype_ext, dim=1, index=nlist_index) - - if self.se_atten.type_one_side: - # (nf*nl*nnei,) -> (nf*nl*nnei, ng) - gg_t = self.type_embd_data[nei_type.view(-1).to(torch.long)] - else: - atype = atype_ext[:, :nloc] - idx_i = torch.tile( - atype.reshape(-1, 1) * ntypes_with_padding, [1, nnei] - ).view(-1) - idx_j = nei_type.view(-1) - idx = (idx_i + idx_j).to(torch.long) - # (nf x nl x nnei) x ng - gg_t = self.type_embd_data[idx] - - # (nf x nl) x nnei x ng - gg_t = gg_t.view(nfnl, nnei, ng) + def _call_compressed( + self, + coord_ext: torch.Tensor, + atype_ext: torch.Tensor, + nlist: torch.Tensor, + ) -> Any: + """Compressed forward for DPA1 descriptor (strip only).""" + nf, nloc, nnei, ng, nfnl, rr, ss, sw, nlist_masked, type_embedding = _env_mat( + self, coord_ext, atype_ext, nlist + ) + tebd_idx = _strip_pair_index( + self, atype_ext, nlist_masked, type_embedding, nf, nloc, nnei + ) + # (nf x nl) x nnei x ng, gathered from the precomputed type-pair table. + gg_t = self.type_embd_data[tebd_idx].view(nfnl, nnei, ng) if self.se_atten.smooth: gg_t = gg_t * sw.view(nfnl, self.se_atten.nnei, 1) if self.geo_compress: - # Flatten for tabulate op - ss_flat = ss.reshape(-1, 1) - gg_t_flat = gg_t.reshape(-1, gg_t.size(-1)) is_sorted = len(self.se_atten.exclude_types) == 0 xyz_scatter = torch.ops.deepmd.tabulate_fusion_se_atten( self.compress_data[0].contiguous(), self.compress_info[0].cpu().contiguous(), - ss_flat.contiguous(), + ss.reshape(-1, 1).contiguous(), rr.contiguous(), - gg_t_flat.contiguous(), + gg_t.reshape(-1, gg_t.size(-1)).contiguous(), self.se_atten.neuron[-1], is_sorted, )[0] else: - # No geometric compression, run embedding net + attention - # nfnl x nnei x ng + # No geometric compression: run embedding net + attention. gg_s = self.se_atten.embeddings[0].call(ss) - # nfnl x nnei x ng gg = gg_s * gg_t + gg_s + nlist_mask = ( + self.se_atten.emask.build_type_exclude_mask(nlist, atype_ext) + .view(nfnl, nnei) + .to(torch.bool) + ) input_r = torch.nn.functional.normalize( rr.view(-1, self.se_atten.nnei, 4)[:, :, 1:4], dim=-1 ) gg = self.se_atten.dpa1_attention(gg, nlist_mask, input_r=input_r, sw=sw) - # nfnl x 4 x ng xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) - xyz_scatter = xyz_scatter / self.se_atten.nnei - # nfnl x ng x 4 - xyz_scatter_1 = xyz_scatter.permute(0, 2, 1) - # nfnl x ng x 3 - rot_mat = xyz_scatter_1[:, :, 1:4] - # nfnl x 4 x axis_neuron - xyz_scatter_2 = xyz_scatter[:, :, 0 : self.se_atten.axis_neuron] - # nfnl x ng x axis_neuron - result = torch.matmul(xyz_scatter_1, xyz_scatter_2) - # nf x nloc x (ng x axis_neuron) - result = result.view(nf, nloc, ng * self.se_atten.axis_neuron) - # nf x nloc x ng x 3 - rot_mat = rot_mat.view(nf, nloc, ng, 3) - - # Concat type embedding at output if needed + return _grrg_from_moment( + self, + xyz_scatter, + type_embedding, + coord_ext, + atype_ext, + nf, + nloc, + nnei, + ng, + sw, + ) + + def _call_graph_compress_reference( + self, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor, + ) -> Any: + """Reference geo-compressed strip descriptor on the edge stream (attn-free). + + The ``DP_CUDA_INFER == 0`` graph path and the CPU / trace-time fallback. + Evaluates the tabulated geometric embedding with the original fused + table operator ``deepmd::tabulate_fusion_se_atten``, treating each edge + as a one-neighbor block (``nloc = E``, ``nnei = 1``) so the operator + returns the per-edge moment outer product ``(E, 4, ng)``; a + ``segment_sum`` over edge centers then forms the per-node moment. This + matches the dense :meth:`DescrptDPA1._call_compressed` (same table, same + gate) to the fp32 summation-order floor, and composes under autograd so + the force / virial assembly differentiates through it. Not + ``make_fx``-traceable (the table operator has no meta kernel); the + exportable graph is the level >= 1 CUDA kernel. + + Parameters + ---------- + graph : NeighborGraph + Lowered neighbor graph; ``edge_vec`` is in the model precision. + atype : torch.Tensor + Flat node atom types with shape (N,), int64. + type_embedding : torch.Tensor + Type embedding table with shape (ntypes + 1, tebd_dim). + + Returns + ------- + grrg : torch.Tensor + Descriptor with shape (N, ng * axis [+ tebd_dim]). + rot_mat : torch.Tensor + Equivariant rotation matrix with shape (N, ng, 3). + """ + se = self.se_atten + ng = int(se.neuron[-1]) + n_total = atype.shape[0] + src, dst = graph.edge_index[0], graph.edge_index[1] + center_type, nei_type = atype[dst], atype[src] + + # === Step 1. Per-edge environment matrix (canonical statistics slot) === + ev = graph.edge_vec + length = ev.norm(dim=-1, keepdim=True) + # Padding edges enter the switch at |r| + 1 so the 1/q factors stay + # finite; their moment weight is zeroed by the edge mask below. + length = length + (~graph.edge_mask[:, None]).to(length.dtype) + q = length + se.env_protection + u = ((length - se.rcut_smth) / (se.rcut - se.rcut_smth)).clamp(0.0, 1.0) + sw = u**3 * (-6 * u**2 + 15 * u - 10) + 1.0 + em = torch.cat([sw / q, ev * (sw / q**2)], dim=-1) + rr = (em - se.mean[:, 0, :][center_type]) / se.stddev[:, 0, :][center_type] + + # === Step 2. Strip type-pair gate from the precomputed table === + ntypes = type_embedding.shape[0] + pair_idx = nei_type if se.type_one_side else center_type * ntypes + nei_type + gate = self.type_embd_data[pair_idx] + if se.smooth: + gate = gate * sw + + # === Step 3. Tabulated embedding and per-edge moment outer product === + # Each edge is a one-neighbor block, so the operator's neighbor sum is + # the identity and it returns the per-edge outer product em (x) g. + is_sorted = len(se.exclude_types) == 0 + outer = torch.ops.deepmd.tabulate_fusion_se_atten( + self.compress_data[0].contiguous(), + self.compress_info[0].cpu().contiguous(), + rr[:, 0:1].contiguous(), + rr.reshape(-1, 1, 4).contiguous(), + gate.contiguous(), + ng, + is_sorted, + )[0] + + # === Step 4. Moment reduction and G^T G contraction === + outer = outer * graph.edge_mask[:, None, None].to(outer.dtype) + gr = torch.zeros(n_total, 4, ng, dtype=outer.dtype, device=outer.device) + gr.index_add_(0, dst, outer) + gr = gr / se.nnei + gr_perm = gr.permute(0, 2, 1) # (N, ng, 4) + rot_mat = gr_perm[:, :, 1:4] + gr_sub = gr[:, :, : se.axis_neuron] # (N, 4, axis) + grrg = torch.matmul(gr_perm, gr_sub).reshape(n_total, ng * se.axis_neuron) + grrg = grrg.to(graph.edge_vec.dtype) if self.concat_output_tebd: - nall = coord_ext.view(nf, -1).shape[1] // 3 - atype_embd_ext = type_embedding[atype_ext.reshape(-1)].view( - nf, nall, self.tebd_dim + grrg = torch.cat([grrg, type_embedding[atype]], dim=-1) + return grrg, rot_mat + + def _call_graph_cuda_compress( + self, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor, + ) -> Any: + """Fused CUDA graph-native geo-compressed strip descriptor (attn-free). + + Numerically equivalent to :meth:`DescrptDPA1._call_compressed` through + the :func:`~deepmd.kernels.cuda.dpa1.graph_compress.dpa1_graph_compress` + operator: the environment matrix, quintic table lookup, strip type-pair + gate, moment reduction and ``G^T G`` contraction collapse into one CUDA + mega kernel whose registered backward exposes the ``edge_vec`` gradient + for the analytic force / virial assembly. + """ + return dpa1_graph_compress(self, graph, atype, type_embedding) + + def _call_graph_cuda( + self, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor, + ) -> Any: + """Fused CUDA graph-native descriptor (concat, attn-free). + + Numerically equivalent to :meth:`DescrptDPA1DP.call_graph` through the + :func:`~deepmd.kernels.cuda.dpa1.graph_descriptor.dpa1_graph_descriptor` + operator: the environment matrix, embedding MLP, moment reduction and + ``G^T G`` contraction collapse into one CUDA mega kernel whose + registered backward exposes the ``edge_vec`` gradient for the analytic + force / virial assembly. Composes under ``make_fx`` / ``torch.export`` + so the operator is baked into the pt_expt graph-form ``.pt2``. + """ + return dpa1_graph_descriptor(self, graph, atype, type_embedding) + + def fused_energy_force_graph( + self, + fit: Any, + graph: Any, + atype: torch.Tensor, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + do_atomic_virial: bool, + ) -> 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 + (: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`. + + Parameters + ---------- + fit : EnergyFittingNet + The energy fitting module fused with this descriptor. + graph : NeighborGraph + The lowered neighbor graph (``edge_vec``, ``edge_index``, + ``edge_mask``, ``n_node``). + atype : torch.Tensor + Flat node atom types with shape (N,), int64. + ownership : torch.Tensor + Energy-contributing node mask with shape (N,), bool. + atom_bias : torch.Tensor + Combined fitting and atomic-model bias with shape (ntypes,). + do_atomic_virial : bool + Whether to also assemble the per-atom virial. + + Returns + ------- + tuple[torch.Tensor, ...] or None + ``(energy, atom_energy, force, virial, atom_virial)``, or ``None``. + """ + from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + ) + + if not ( + self._fused_eligible("cuda") + and fitting_eligible(fit) + and _has_graph_fields(graph, _DUAL_CSR_FIELDS) + ): + return None + type_embedding = self.type_embedding.call() + node_capacity = atype.shape[0] + if self.geo_compress: + from deepmd.kernels.cuda.dpa1.graph_compress import ( + dpa1_graph_compress_energy_force, + ef_op_available, + mega_eligible, + ) + + if not (ef_op_available() and mega_eligible(self)): + return None + return dpa1_graph_compress_energy_force( + self, + fit, + graph, + atype, + type_embedding, + ownership, + atom_bias, + node_capacity=node_capacity, + do_atomic_virial=do_atomic_virial, ) - atype_embd = atype_embd_ext[:, :nloc, :] - result = torch.cat( - [result, atype_embd.view(nf, nloc, self.tebd_dim)], dim=-1 + 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( + self, + graph: Any, + atype: torch.Tensor, + type_embedding: torch.Tensor, + ) -> Any: + """Fused graph-native environment convolution (concat / strip, attn-free). + + Bit-exact analogue of :meth:`DescrptDPA1DP.call_graph`: builds the + per-edge environment matrix and embedding input, runs the two embedding + GEMMs on cuBLAS, then folds the last layer, the type-pair gate (strip), + the edge mask, the outer product and the ``segment_sum`` scatter into + :func:`edge_conv`. The two tebd-input modes differ only in the embedding + input (concat folds the type feature in; strip runs on the radial channel + alone) and the gate (strip multiplies by the type-pair table, concat does + not). Composes under ``make_fx`` / ``torch.export`` so the operator is + baked into the pt_expt graph-form ``.pt2``. + """ + se = self.se_atten + strip = se.tebd_input_mode == "strip" + n_total = atype.shape[0] + src = graph.edge_index[0, :] + dst = graph.edge_index[1, :] + center_type = atype[dst] + nei_type = atype[src] + # Per-edge env-mat 4-vector, normalized by the center (dst) atom type; + # mean/stddev are slot-independent, so slot 0 is the canonical vector. + # The fused operator is captured opaquely under the pt_expt trace and + # resolves to the Triton kernel at CUDA runtime (identical to the + # array-API ``edge_env_mat`` reference it routes to at level 0). Strip + # also needs the per-edge switch for the type-pair gate (differentiable + # in ``edge_vec``); the same operator emits it when ``return_sw`` is set. + rr, sw_e = _edge_env_mat_triton( + graph.edge_vec, + center_type, + se.mean[:, 0, :], + se.stddev[:, 0, :], + se.rcut, + se.rcut_smth, + protection=se.env_protection, + edge_mask=graph.edge_mask, + return_sw=True, + ) # (E, 4), (E, 1) + ss = rr[:, 0:1] + ng = se.neuron[-1] + if strip: + # Strip: the geometric net runs on the radial channel alone; the + # type feature enters as the multiplicative gate ``1 + tt[idx] * sw`` + # applied inside edge_conv. ``tt`` is the type-pair table (built like + # the dense strip reference); ``idx`` folds the (center, neighbor) + # pair (two-side) or the neighbor type (one-side); ``sw`` is the + # switch when smooth, else ones (matching the dense path). + emb_in = ss + tt = _type_pair_table(self, type_embedding) + ntypes_with_padding = type_embedding.shape[0] + if se.type_one_side: + gate_idx = nei_type.to(torch.long) + else: + gate_idx = center_type.to( + torch.long + ) * ntypes_with_padding + nei_type.to(torch.long) + sw_gate = ( + sw_e.reshape(-1) + if se.smooth + else torch.ones(rr.shape[0], dtype=rr.dtype, device=rr.device) ) + gated = 1 + else: + # concat embedding input: radial channel plus the neighbor (and, two- + # side, center) type embeddings. Ghost type == owner type, so + # gathering by the local owner reproduces the dense neighbor tebd. + nlist_tebd = type_embedding[nei_type] # (E, tebd_dim) + if se.type_one_side: + emb_in = torch.cat([ss, nlist_tebd], dim=-1) + else: + center_tebd = type_embedding[center_type] # (E, tebd_dim) + emb_in = torch.cat([ss, nlist_tebd, center_tebd], dim=-1) + tt, gate_idx, sw_gate = edge_concat_gate_placeholders(rr, ng) + gated = 0 - return result, rot_mat, None, None, sw.view(nf, nloc, nnei, 1) + # Embedding net through the penultimate layer (cuBLAS); fold the final + # layer's pre-activation into edge_conv. + *head, last = se.embeddings[0].layers + h = emb_in + for layer in head: + h = layer.call(h) + z2 = embed_last_gemm(h, last.w, last.b) + h1_dim, out_dim = last.w.shape + if last.resnet and out_dim == 2 * h1_dim: + resnet_mult = 2 + elif last.resnet and out_dim == h1_dim: + resnet_mult = 1 + else: + resnet_mult = 0 + idt = ( + last.idt + if last.idt is not None + else torch.ones(out_dim, dtype=z2.dtype, device=z2.device) + ) + act = ACT_CODES[str(last.activation_function).lower()] + # Unnormalized moment (N, 4, ng); apply the 1 / nnei normalization here. + gr = ( + edge_conv( + z2.contiguous(), + h.contiguous(), + idt, + tt.contiguous(), + gate_idx.contiguous(), + sw_gate.contiguous(), + rr.contiguous(), + dst, + graph.edge_mask, + n_total, + resnet_mult, + act, + gated, + ) + / se.nnei + ) + ng = se.neuron[-1] + # G^T G contraction over the 4-axis (mirrors the dense _grrg_from_moment + # on the (N, 4, ng) moment): permute to (N, ng, 4), contract the first + # ``axis_neuron`` channels. + gr_perm = gr.permute(0, 2, 1) # (N, ng, 4) + rot_mat = gr_perm[:, :, 1:4] # (N, ng, 3) + gr_sub = gr[:, :, : se.axis_neuron] # (N, 4, axis) + grrg = torch.matmul(gr_perm, gr_sub) # (N, ng, axis) + grrg = grrg.reshape(n_total, ng * se.axis_neuron).to(graph.edge_vec.dtype) + if self.concat_output_tebd: + atype_embd = type_embedding[atype] # (N, tebd_dim) + grrg = torch.cat([grrg, atype_embd], dim=-1) + return grrg, rot_mat diff --git a/deepmd/pt_expt/descriptor/se_atten_v2.py b/deepmd/pt_expt/descriptor/se_atten_v2.py index e0eb3acac3..c50bf8292d 100644 --- a/deepmd/pt_expt/descriptor/se_atten_v2.py +++ b/deepmd/pt_expt/descriptor/se_atten_v2.py @@ -57,6 +57,65 @@ def call(self, *args: Any, **kwargs: Any) -> Any: return DescrptDPA1.call(self, *args, **kwargs) + def _fused_eligible(self, *args: Any, **kwargs: Any) -> bool: + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + return DescrptDPA1._fused_eligible(self, *args, **kwargs) + + def fused_energy_force_graph(self, *args: Any, **kwargs: Any) -> Any: + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + return DescrptDPA1.fused_energy_force_graph(self, *args, **kwargs) + + # Graph-lower routing and the fused-kernel helpers it dispatches to, + # forwarded to DescrptDPA1 so the accelerated CUDA / Triton / tabulated + # graph paths serve se_atten_v2 rather than only the dpmodel reference. + def call_graph(self, *args: Any, **kwargs: Any) -> Any: + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + return DescrptDPA1.call_graph(self, *args, **kwargs) + + def _call_graph_cuda(self, *args: Any, **kwargs: Any) -> Any: + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + return DescrptDPA1._call_graph_cuda(self, *args, **kwargs) + + def _call_graph_cuda_compress(self, *args: Any, **kwargs: Any) -> Any: + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + return DescrptDPA1._call_graph_cuda_compress(self, *args, **kwargs) + + def _call_graph_triton(self, *args: Any, **kwargs: Any) -> Any: + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + return DescrptDPA1._call_graph_triton(self, *args, **kwargs) + + def _call_graph_compress_reference(self, *args: Any, **kwargs: Any) -> Any: + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + return DescrptDPA1._call_graph_compress_reference(self, *args, **kwargs) + + def _call_triton(self, *args: Any, **kwargs: Any) -> Any: + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + return DescrptDPA1._call_triton(self, *args, **kwargs) + def _call_compressed(self, *args: Any, **kwargs: Any) -> Any: from deepmd.pt_expt.descriptor.dpa1 import ( DescrptDPA1, diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 83417c68f2..3d9bb6506a 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -106,22 +106,38 @@ def enable_compression( # 4. Serialize the compressed model dict (includes tabulated data) compressed_model_dict = model.serialize() - # 5. Re-export: trace an UNCOMPRESSED model for the exported program - # (make_fx cannot trace through custom tabulate ops), but store - # the full compressed dict in model.json so that deserialize() - # restores compression state. - log.info("Re-exporting compressed model...") - uncompressed_data = { - "model": model_dict["model"], # original uncompressed model - "model_def_script": model_dict.get("model_def_script"), - } - deserialize_to_file( - output, - uncompressed_data, - model_json_override={ - "model": compressed_model_dict, - "model_def_script": model_dict.get("model_def_script"), - "min_nbor_dist": float(min_nbor_dist), - }, + # 5. Re-export the compressed model. + # + # A geometrically compressed graph-lower descriptor (DPA1 / se_atten strip, + # attn_layer == 0) evaluates its tabulated embedding through fused CUDA + # operators that make_fx can trace, so its compressed graph exports directly + # to a graph-lower ``.pt2``: ``deserialize_to_file`` bakes the end-to-end + # fused table operator and the mandatory per-atom virial (see its docstring). + # Every other compressed descriptor keeps tabulated operators make_fx cannot + # trace: those export the UNCOMPRESSED graph and carry the compressed dict in + # ``model.json`` so ``deserialize()`` restores the compression state for the + # Python inference path. + from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower, ) + + model_def_script = model_dict.get("model_def_script") + if output.endswith(".pt2") and model_uses_graph_lower(model): + log.info("Re-exporting compressed graph...") + deserialize_to_file( + output, + {"model": compressed_model_dict, "model_def_script": model_def_script}, + lower_kind="auto", + ) + else: + log.info("Re-exporting compressed model...") + deserialize_to_file( + output, + {"model": model_dict["model"], "model_def_script": model_def_script}, + model_json_override={ + "model": compressed_model_dict, + "model_def_script": model_def_script, + "min_nbor_dist": float(min_nbor_dist), + }, + ) log.info("Compressed model saved to %s", output) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 3367ee4579..e49a73a518 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -576,11 +576,11 @@ def freeze( # scatter off the single shared backward). do_atomic_virial = False if lower_kind == "graph": - from deepmd.pt_expt.train.training import ( - _model_uses_graph_lower, + from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower, ) - if not _model_uses_graph_lower(m): + if not model_uses_graph_lower(m): raise ValueError( "lower_kind='graph' requires a graph-eligible model " "(mixed_types and a descriptor exposing uses_graph_lower()==True, " diff --git a/deepmd/pt_expt/fitting/ener_fitting.py b/deepmd/pt_expt/fitting/ener_fitting.py index e931c72b5c..f618095141 100644 --- a/deepmd/pt_expt/fitting/ener_fitting.py +++ b/deepmd/pt_expt/fitting/ener_fitting.py @@ -3,7 +3,17 @@ Any, ) +import torch + from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnergyFittingNetDP +from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + graph_fitting, +) +from deepmd.kernels.cuda.graph_fitting import op_available as cuda_fitting_available +from deepmd.kernels.utils import ( + cuda_infer_level, +) from deepmd.pt_expt.common import ( torch_module, ) @@ -22,3 +32,42 @@ def share_params(self, *args: Any, **kwargs: Any) -> None: ) return InvarFitting.share_params(self, *args, **kwargs) + + def call_graph( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Graph-native fitting forward, fused on CUDA when eligible. + + At ``DP_CUDA_INFER >= 1`` an inference-mode call on an eligible + network (see + :func:`~deepmd.kernels.cuda.graph_fitting.fitting_eligible`) + routes through the fused cuBLAS operator; anything else keeps the + dpmodel reference. Routing is device-free, so a CPU ``make_fx`` trace + bakes the operator into the exported graph. + """ + if ( + not self.training + and fparam is None + and aparam is None + and cuda_infer_level() >= 1 + and cuda_fitting_available() + and fitting_eligible(self) + ): + return graph_fitting(self, descriptor, atype) + return EnergyFittingNetDP.call_graph( + self, + descriptor, + atype, + gr=gr, + g2=g2, + h2=h2, + fparam=fparam, + aparam=aparam, + ) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index f198d4bbe4..66efeddb89 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -187,10 +187,9 @@ def __init__( # neighbor_graph_method is consumed ONLY by graph-form .pt2 eval # (_eval_model_graph); fail fast instead of silently ignoring it on # nlist-form artifacts (there, the builder knob is nlist_backend). - if ( - neighbor_graph_method != "dense" - and getattr(self, "metadata", {}).get("lower_input_kind") != "graph" - ): + if neighbor_graph_method != "dense" and getattr(self, "metadata", {}).get( + "lower_input_kind" + ) not in ("graph", "dpa1_canonical"): raise ValueError( f"neighbor_graph_method={neighbor_graph_method!r} only applies to " "graph-form .pt2 artifacts (lower_input_kind == 'graph'); this " @@ -1527,7 +1526,7 @@ def _eval_model( 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", "dpa1_canonical"): return self._eval_model_graph( coords, cells, atom_types, fparam, aparam, request_defs, charge_spin ) @@ -1767,11 +1766,14 @@ def _eval_model_graph( Builds a carry-all :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph` from the eval system at its exact (tight) edge count and feeds the positional schema - ``(atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, - charge_spin)`` to the exported forward. The AOTI artifact's edge axis - is DYNAMIC (B2.0), so no ``edge_capacity`` padding is needed. The - forward returns the LOCAL public keys directly, so results are reshaped - without ``communicate_extended_output``. + ``(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)`` to the exported forward. The AOTI + artifact's edge axis is dynamic, so no ``edge_capacity`` padding is needed. The + ``graph_edge_dtype`` metadata selects float32 geometry for compressed + DPA1 and float64 for generic graph descriptors. The forward returns the + LOCAL public keys directly, so results are reshaped without + ``communicate_extended_output``. """ from deepmd.pt_expt.utils.env import ( DEVICE, @@ -1786,8 +1788,8 @@ def _eval_model_graph( coord_input = coords.reshape(nframes, natoms, 3) box_input = cells.reshape(nframes, 9) if cells is not None else None - # Dynamic edge axis (B2.0): build the carry-all graph at its exact edge - # count (no static padding); the AOTI artifact accepts any E. + # Build the carry-all graph at its exact edge count; the exported edge + # axis is dynamic. graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE) atype_t = torch.tensor( @@ -1799,24 +1801,87 @@ def _eval_model_graph( edge_index_t = torch.as_tensor( graph.edge_index, dtype=torch.int64, device=DEVICE ) - edge_vec_t = torch.as_tensor(graph.edge_vec, dtype=torch.float64, device=DEVICE) + edge_dtype = ( + torch.float32 + if self.metadata.get("graph_edge_dtype") == "float32" + else torch.float64 + ) + edge_vec_t = torch.as_tensor( + graph.edge_vec, + dtype=edge_dtype, + device=DEVICE, + ) edge_mask_t = torch.as_tensor(graph.edge_mask, dtype=torch.bool, device=DEVICE) - - fparam_t, aparam_t = self._prepare_optional_lower_inputs( - fparam, aparam, nframes, natoms, DEVICE + destination_order_t = torch.as_tensor( + graph.destination_order, + device=DEVICE, ) - charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) - - model_inputs = ( - atype_t, - n_node_t, - edge_index_t, - edge_vec_t, - edge_mask_t, - fparam_t, - aparam_t, - charge_spin_t, + destination_row_ptr_t = torch.as_tensor( + graph.destination_row_ptr, + dtype=torch.int64, + device=DEVICE, + ) + source_order_t = torch.as_tensor( + graph.source_order, + device=DEVICE, ) + source_row_ptr_t = torch.as_tensor( + graph.source_row_ptr, + dtype=torch.int64, + device=DEVICE, + ) + + if self.metadata.get("lower_input_kind") == "dpa1_canonical": + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + ) + + generic_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), + ) + compact = canonical_graph_from_neighbor_graph(generic_graph) + 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, + ) + else: + fparam_t, aparam_t = self._prepare_optional_lower_inputs( + fparam, aparam, nframes, natoms, DEVICE + ) + charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) + 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, + fparam_t, + aparam_t, + charge_spin_t, + ) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: @@ -1847,7 +1912,9 @@ def _build_eval_graph( Dispatches on ``self._neighbor_graph_method``: ``dense``/``ase`` run backend-agnostic (numpy); ``vesin``/``nv`` run on-device (torch, O(N)). All backends emit the SAME neighbor set (carry-all, sel-free), so the - selection is a pure performance choice and results are unchanged. + selection is a pure performance choice and results are unchanged. The + result is canonicalized to the destination-major graph-form ``.pt2`` + ABI after construction. """ method = self._neighbor_graph_method # Model-level ``pair_exclude_types`` is a graph-BUILD transform @@ -1861,7 +1928,12 @@ def _build_eval_graph( ) return build_neighbor_graph( - coord_input, atom_types, box_input, self._rcut, pair_excl=pair_excl + coord_input, + atom_types, + box_input, + self._rcut, + canonicalize=True, + pair_excl=pair_excl, ) if method == "ase": from deepmd.dpmodel.utils.neighbor_graph import ( @@ -1869,7 +1941,12 @@ def _build_eval_graph( ) return build_neighbor_graph_ase( - coord_input, atom_types, box_input, self._rcut, pair_excl=pair_excl + coord_input, + atom_types, + box_input, + self._rcut, + canonicalize=True, + pair_excl=pair_excl, ) if method in ("vesin", "nv"): cc = torch.as_tensor(coord_input, dtype=torch.float64, device=device) @@ -1887,13 +1964,25 @@ def _build_eval_graph( ) return build_neighbor_graph_vesin( - cc, aa, bb, self._rcut, pair_excl=pair_excl + cc, + aa, + bb, + self._rcut, + canonicalize=True, + pair_excl=pair_excl, ) from deepmd.pt_expt.utils.nv_graph_builder import ( build_neighbor_graph_nv, ) - return build_neighbor_graph_nv(cc, aa, bb, self._rcut, pair_excl=pair_excl) + return build_neighbor_graph_nv( + cc, + aa, + bb, + self._rcut, + canonicalize=True, + pair_excl=pair_excl, + ) raise ValueError( f"unknown neighbor_graph_method {method!r}; " "use 'dense', 'ase', 'vesin', or 'nv'" diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 653f323404..24c469c0a5 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -18,6 +18,15 @@ frame_id_from_n_node, segment_sum, ) +from deepmd.kernels.cuda.edge_force_virial import ( + edge_force_virial as fused_edge_force_virial, +) +from deepmd.kernels.cuda.edge_force_virial import ( + op_available as fused_scatter_available, +) +from deepmd.kernels.utils import ( + cuda_infer_level, +) from deepmd.pt.utils import ( env, ) @@ -29,15 +38,33 @@ def edge_energy_deriv( edge_index: torch.Tensor, edge_mask: torch.Tensor, n_node: torch.Tensor, + destination_order: torch.Tensor | None = None, + destination_row_ptr: torch.Tensor | None = None, + source_order: torch.Tensor | None = None, + source_row_ptr: torch.Tensor | None = None, node_capacity: int | None = None, *, do_atomic_virial: bool = False, create_graph: bool = False, + force_precision: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """Return (force, atom_virial_or_None, virial) from a graph energy. g_e = dE/d(edge_vec) via one torch.autograd.grad, then the shared - edge_force_virial scatter. + edge_force_virial scatter. At ``DP_CUDA_INFER >= 1`` (and with the C++ + operator library loaded) the scatter runs through the fused + ``deepmd::edge_force_virial`` operator instead of the array-API kernel + chain; the routing is device-free so a CPU ``make_fx`` trace bakes the + operator into the exported graph. + + ``g_e`` is the gradient with respect to the fp64 ``edge_vec`` leaf, but the + descriptor and fitting compute -- hence g_e's numerical content -- are in + the model precision (fp32 for an fp32 model). For inference, scattering the + force / virial in that ``force_precision`` rather than the fp64 leaf dtype + halves the atomic-scatter traffic at no accuracy cost: the per-node force is + a small neighbor sum and the per-frame virial uses a hierarchical block + reduction. A retained graph (training / double backward) keeps the fp64 leaf + dtype so training numerics are unchanged. Parameters ---------- @@ -51,6 +78,10 @@ def edge_energy_deriv( (E,) valid-edge mask. n_node (nf,) per-frame node counts. + destination_order, source_order + (E,) destination/source-grouped edge permutations. + destination_row_ptr, source_row_ptr + (N + 1,) destination/source CSR offsets. node_capacity Static node-axis size ``N``. ``None`` (eager default) falls back to ``int(n_node.sum())``. Pass a static value (e.g. ``atype.shape[0]``) @@ -75,9 +106,40 @@ def edge_energy_deriv( create_graph=create_graph, retain_graph=True, ) - force, atom_virial, virial = edge_force_virial( - g_e, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity - ) + if ( + force_precision is not None + and not create_graph + and g_e.dtype != force_precision + ): + g_e = g_e.to(force_precision) + edge_vec = edge_vec.to(force_precision) + if ( + cuda_infer_level() >= 1 + and not create_graph + and fused_scatter_available() + and destination_order is not None + and destination_row_ptr is not None + and source_order is not None + 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( + g_e, + edge_vec, + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + n_node, + n_cap, + do_atomic_virial, + ) + else: + force, atom_virial, virial = edge_force_virial( + g_e, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity + ) return force, (atom_virial if do_atomic_virial else None), virial @@ -89,6 +151,7 @@ def fit_output_to_model_output_graph( create_graph: bool = True, mask: torch.Tensor | None = None, node_capacity: int | None = None, + force_precision: torch.dtype | None = None, ) -> dict[str, torch.Tensor]: """Graph analogue of the dense pt_expt ``fit_output_to_model_output``. @@ -131,6 +194,10 @@ def fit_output_to_model_output_graph( input node axis rather than a re-derived shape -- hardening; the actual CUDA out-of-bounds device-assert is prevented by the index clamp in :func:`~deepmd.dpmodel.utils.neighbor_graph.derivatives.edge_force_virial`. + force_precision + Compute precision (model dtype) in which to assemble the force / virial + during inference, decoupled from the fp64 ``edge_vec`` leaf; see + :func:`edge_energy_deriv`. ``None`` keeps the leaf dtype. Returns ------- @@ -201,9 +268,14 @@ def fit_output_to_model_output_graph( edge_index, edge_mask, n_node, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, node_capacity=N, do_atomic_virial=(vdef.c_differentiable and do_atomic_virial), create_graph=create_graph, + force_precision=force_precision if not create_graph else None, ) # force (N, 3) -> (N, 1, 3) [flat; caller unravels at I/O boundary] ff_list.append(force.reshape(N, 1, 3)) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 140141dd5e..d0990ffe17 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -90,6 +90,168 @@ def enable_hessian(self) -> None: self.requires_hessian("energy") self._hessian_enabled = True + def forward_lower_canonical_graph( + 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, + *, + do_atomic_virial: bool, + ) -> dict[str, torch.Tensor]: + """Evaluate an eligible compressed DPA1 deployment graph. + + Parameters + ---------- + atype + Flat local-plus-halo atom types with shape ``(N,)``, int64. + n_node + Per-frame total node counts with shape ``(nf,)``, int64. + n_local + Per-frame owned node counts with shape ``(nf,)``, int64. + source + Source-node indices with shape ``(S,)``, int32 or int64. + edge_vec + Destination-major edge vectors with shape ``(S, 3)``, float32. + destination_row_ptr + Destination CSR offsets with shape ``(N + 1,)``, int64. + source_row_ptr + Source CSR offsets with shape ``(N + 1,)``, int64. + source_order + Source-grouped edge positions with shape ``(S,)`` and the same + dtype as ``source``. + do_atomic_virial + Whether to return the per-node virial. + + Returns + ------- + dict[str, torch.Tensor] + Public energy-model outputs on the flat node axis. + """ + from deepmd.kernels.cuda.dpa1.canonical import ( + canonical_model_eligible, + dpa1_canonical_compress_energy_force, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + DPA1CanonicalGraph, + validate_canonical_graph_shapes, + ) + + if not canonical_model_eligible(self): + raise ValueError("model is not eligible for compact canonical deployment") + graph = DPA1CanonicalGraph( + n_node=n_node, + n_local=n_local, + source=source, + edge_vec=edge_vec, + destination_row_ptr=destination_row_ptr, + source_row_ptr=source_row_ptr, + source_order=source_order, + ) + validate_canonical_graph_shapes(graph, atype.shape[0]) + atype, output_mask = self.atomic_model._prepare_graph_nodes( + n_node, + n_local, + atype, + edge_vec, + ) + 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.type_embedding.call(), + output_mask, + atom_bias, + do_atomic_virial, + ) + ) + result = { + "atom_energy": atom_energy, + "energy": energy, + "force": force, + "virial": virial, + "mask": output_mask.to(torch.int32), + } + if do_atomic_virial: + result["atom_virial"] = atom_virial + return result + + 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, + *, + do_atomic_virial: bool, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + """Trace the compact canonical deployment forward with ``make_fx``. + + 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. + 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 eight-input compact deployment module. + """ + 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, + ) -> dict[str, torch.Tensor]: + return 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, + ) + + return make_fx(fn, **make_fx_kwargs)( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + ) + def forward( self, coord: torch.Tensor, @@ -280,13 +442,19 @@ def forward_lower_graph_exportable( self, atype: torch.Tensor, n_node: torch.Tensor, + n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, + destination_sorted: bool = False, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_common_lower_graph`` into an exportable module with @@ -298,15 +466,26 @@ def forward_lower_graph_exportable( Parameters ---------- atype - (N,) flat local atom types, ``N == sum(n_node)``. + (N,) flat local-plus-halo atom types, ``N == sum(n_node)``. n_node - (nf,) per-frame local atom counts. + (nf,) per-frame total node counts. + n_local + (nf,) per-frame owned node counts. edge_index (2, E) ``[src, dst]`` edge endpoints (flat local indices). edge_vec (E, 3) neighbor-minus-center edge vectors (sample for tracing). edge_mask (E,) valid-edge mask (sample for tracing). + destination_order + (E,) destination-grouped edge permutation. + source_order + (E,) source-grouped edge permutation. + destination_row_ptr, source_row_ptr + (N + 1,) destination/source CSR offsets. + destination_sorted + Static export-time assertion that the payload is destination-major + and ``destination_order`` is identity. fparam, aparam, do_atomic_virial, charge_spin As in ``forward_lower``. **make_fx_kwargs @@ -317,26 +496,32 @@ def forward_lower_graph_exportable( ------- torch.nn.Module A traced module whose ``forward`` accepts - ``(atype, n_node, edge_index, edge_vec, edge_mask, - fparam, aparam, charge_spin)`` and returns a dict with the - public keys: ``atom_energy``, ``energy``, ``force``, + ``(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)`` and returns a dict + with the public keys: ``atom_energy``, ``energy``, ``force``, ``virial``, ``atom_virial`` (the last only when - ``do_atomic_virial``). Unlike the dense + ``do_atomic_virial``). Unlike the dense :meth:`forward_lower_exportable` (which emits ``extended_force`` / - ``extended_virial`` over the ghost-padded extended region), the - graph path is LOCAL-only (``N == sum(n_node)`` nodes, no ghosts), - so it emits ``force`` / ``atom_virial`` directly. + ``extended_virial``), the graph path emits ``force`` and + ``atom_virial`` directly on the local-plus-halo node axis. """ traced = self.forward_common_lower_graph_exportable( atype, n_node, + n_local, edge_index, edge_vec, edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, fparam=fparam, aparam=aparam, charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, + destination_sorted=destination_sorted, **make_fx_kwargs, ) @@ -348,9 +533,14 @@ def forward_lower_graph_exportable( def fn( atype: torch.Tensor, n_node: torch.Tensor, + n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, @@ -358,9 +548,14 @@ def fn( model_ret = traced( 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, @@ -374,5 +569,17 @@ def fn( ) return make_fx(fn, **make_fx_kwargs)( - atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin + 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, ) diff --git a/deepmd/pt_expt/model/graph_lower.py b/deepmd/pt_expt/model/graph_lower.py new file mode 100644 index 0000000000..fce1a63814 --- /dev/null +++ b/deepmd/pt_expt/model/graph_lower.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared graph-lower model capability checks.""" + +from typing import ( + Any, +) + + +def model_uses_graph_lower(model: Any) -> bool: + """Return whether a model's default lower uses ``NeighborGraph``. + + Parameters + ---------- + model : Any + Model exposing the atomic-model and descriptor capability interfaces. + + Returns + ------- + bool + ``True`` when the model is mixed-type and its descriptor enables the + graph lower. + """ + mixed_types = getattr(model, "mixed_types", None) + if mixed_types is None: + return False + try: + if not bool(mixed_types()): + return False + except (AttributeError, NotImplementedError): + return False + + descriptor = getattr(getattr(model, "atomic_model", None), "descriptor", None) + uses_graph_lower = getattr(descriptor, "uses_graph_lower", None) + if uses_graph_lower is None: + return False + try: + return bool(uses_graph_lower()) + except (AttributeError, NotImplementedError): + return False diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index a080600709..9885d599c6 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -16,13 +16,22 @@ from deepmd.dpmodel.atomic_model.base_atomic_model import ( BaseAtomicModel, ) +from deepmd.dpmodel.common import ( + get_xp_precision, +) from deepmd.dpmodel.model.make_model import make_model as make_model_dp from deepmd.dpmodel.output_def import ( OutputVariableDef, ) +from deepmd.kernels.utils import ( + cuda_infer_level, +) from deepmd.pt_expt.common import ( torch_module, ) +from deepmd.pt_expt.utils.graph_csr import ( + validate_graph_csr_for_export, +) from .edge_transform_output import ( fit_output_to_model_output_graph, @@ -32,6 +41,58 @@ ) +def _fused_energy_force_graph( + model: Any, + graph: Any, + atype: torch.Tensor, + do_atomic_virial: bool, +) -> dict[str, torch.Tensor] | None: + """End-to-end energy / force / virial via fused inference operators. + + At ``DP_CUDA_INFER >= 2``, a descriptor that exposes + ``fused_energy_force_graph`` (the DPA1 ``se_atten``, attention-free family) + emits descriptor, fitting, analytic descriptor-backward, and CSR + force/virial custom operators. Force is returned as a value, so the graph + carries no autograd tape. The descriptor owns backend eligibility and kernel + dispatch (embedding-MLP vs tabulated); this routine only assembles the flat + model dict. Returns the same keys as + :func:`fit_output_to_model_output_graph`, or ``None`` when no descriptor + fused path applies -- the caller then uses the level-1 autograd lower. + """ + am = model.atomic_model + desc = getattr(am, "descriptor", None) + fit = getattr(am, "fitting_net", None) + fused = getattr(desc, "fused_energy_force_graph", None) + if fused is None or fit is None: + return None + graph, atype, output_mask = am._prepare_graph_inputs(graph, atype) + atom_bias = fit.bias_atom_e[:, 0] + am.out_bias[0, :, 0] + out = fused( + fit, + graph, + atype, + output_mask, + atom_bias, + do_atomic_virial, + ) + if out is None: + return None + energy, atom_energy, force, virial, atom_virial = out + n = atype.shape[0] + nf = graph.n_node.shape[0] + var = fit.var_name + ret = { + var: atom_energy, + var + "_redu": energy, + var + "_derv_r": force.reshape(n, 1, 3), + var + "_derv_c_redu": virial.reshape(nf, 1, 9), + "mask": output_mask.to(torch.int32), + } + if do_atomic_virial: + ret[var + "_derv_c"] = atom_virial.reshape(n, 1, 9) + return ret + + def _pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor: """Append a single ``-1`` column to ``nlist`` for export-time tracing. @@ -284,9 +345,15 @@ def forward_common_lower_graph( self, atype: torch.Tensor, n_node: torch.Tensor, + n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, + destination_order: torch.Tensor | None = None, + destination_row_ptr: torch.Tensor | None = None, + source_order: torch.Tensor | None = None, + source_row_ptr: torch.Tensor | None = None, + destination_sorted: bool = False, do_atomic_virial: bool = False, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, @@ -314,15 +381,28 @@ def forward_common_lower_graph( Parameters ---------- atype - (N,) flat LOCAL atom types, ``N == sum(n_node)``. + (N,) flat local-plus-halo atom types, ``N == sum(n_node)``. n_node - (nf,) per-frame local atom counts. + (nf,) per-frame total node counts, including halo nodes. + n_local + (nf,) per-frame owned node counts. edge_index (2, E) ``[src, dst]`` edge endpoints (flat local indices). edge_vec (E, 3) neighbor-minus-center edge vectors. edge_mask (E,) valid-edge mask. + destination_order + (E,) destination-grouped edge permutation. + destination_row_ptr + (N + 1,) destination CSR offsets into ``destination_order``. + source_order + (E,) edge permutation grouped by source node. + source_row_ptr + (N + 1,) source CSR offsets into ``source_order``. + destination_sorted + Whether the payload is destination-major and + ``destination_order`` is the identity permutation. do_atomic_virial Whether to also return the per-atom virial ``_derv_c``. fparam @@ -330,8 +410,8 @@ def forward_common_lower_graph( aparam Atomic parameter, ``(nf, nloc, nda)``. charge_spin - charge/spin conditioning. Ignored in PR-A; accepted for ABI - stability with charge/spin-conditioned descriptors. + Charge/spin conditioning. Accepted for descriptors that expose + this input. Returns ------- @@ -353,7 +433,19 @@ def forward_common_lower_graph( edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, + n_local=n_local, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=destination_sorted, ) + # Level 2 emits force as a value through the inference-only custom + # operator pipeline. Ineligible models use the autograd lower. + if not self.training and cuda_infer_level() >= 2: + fused = _fused_energy_force_graph(self, graph, atype, do_atomic_virial) + if fused is not None: + return fused atomic_ret = self.atomic_model.forward_common_atomic_graph( graph, atype, @@ -370,6 +462,13 @@ def forward_common_lower_graph( do_atomic_virial=do_atomic_virial, create_graph=self.training, mask=atomic_ret["mask"] if "mask" in atomic_ret else None, + # Assemble force / virial in the descriptor compute precision + # (fp32 for an fp32 model) rather than the fp64 edge_vec leaf; + # the gradient content is only that precision, so the coarser + # scatter halves the atomic traffic at no accuracy cost. + force_precision=get_xp_precision( + torch, self.atomic_model.descriptor.precision + ), # Bound the per-node scatter by the INPUT node axis (the symbol # ``edge_index`` indexes into), not the re-derived fitting-output # shape -- avoids a CUDA out-of-bounds device-assert under @@ -468,29 +567,56 @@ def _call_common_graph( "graph lower (e.g. dpa1 attn_layer=0)" ) rcut = self.get_rcut() - # Model-level pair_exclude_types is a graph-BUILD transform - # (decision #18): apply it here, at the single owning site, so the - # exported lower (forward_common_atomic_graph, which no longer - # re-applies it) consumes a pre-excluded edge_mask. + with_csr = ( + not self.training + and cuda_infer_level() >= 1 + and bool(getattr(descriptor, "geo_compress", False)) + ) pair_excl = getattr(self.atomic_model, "pair_excl", None) if method == "dense": - ng = build_neighbor_graph(cc, atype, bb, rcut, pair_excl=pair_excl) + ng = build_neighbor_graph( + cc, + atype, + bb, + rcut, + with_csr=with_csr, + pair_excl=pair_excl, + ) elif method == "ase": - ng = build_neighbor_graph_ase(cc, atype, bb, rcut, pair_excl=pair_excl) + ng = build_neighbor_graph_ase( + cc, + atype, + bb, + rcut, + with_csr=with_csr, + pair_excl=pair_excl, + ) elif method == "vesin": from deepmd.pt_expt.utils.vesin_graph_builder import ( build_neighbor_graph_vesin, ) ng = build_neighbor_graph_vesin( - cc, atype, bb, rcut, pair_excl=pair_excl + cc, + atype, + bb, + rcut, + with_csr=with_csr, + pair_excl=pair_excl, ) elif method == "nv": from deepmd.pt_expt.utils.nv_graph_builder import ( build_neighbor_graph_nv, ) - ng = build_neighbor_graph_nv(cc, atype, bb, rcut, pair_excl=pair_excl) + ng = build_neighbor_graph_nv( + cc, + atype, + bb, + rcut, + with_csr=with_csr, + pair_excl=pair_excl, + ) else: raise ValueError( f"unknown neighbor_graph_method {method!r}; " @@ -501,9 +627,15 @@ def _call_common_graph( model_predict = self.forward_common_lower_graph( atype_flat, ng.n_node, + ng.n_node, ng.edge_index, ng.edge_vec, ng.edge_mask, + ng.destination_order, + ng.destination_row_ptr, + ng.source_order, + ng.source_row_ptr, + destination_sorted=ng.destination_sorted, do_atomic_virial=do_atomic_virial, fparam=fp, aparam=ap, @@ -670,13 +802,19 @@ def forward_common_lower_graph_exportable( self, atype: torch.Tensor, n_node: torch.Tensor, + n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, + destination_sorted: bool = False, **make_fx_kwargs: Any, ) -> torch.nn.Module: """make_fx trace of ``forward_common_lower_graph`` with ``edge_vec`` @@ -685,15 +823,27 @@ def forward_common_lower_graph_exportable( Parameters ---------- atype - (N,) flat local atom types, ``N == sum(n_node)``. + (N,) flat local-plus-halo atom types, ``N == sum(n_node)``. n_node - (nf,) per-frame local atom counts. + (nf,) per-frame total node counts. + n_local + (nf,) per-frame owned node counts. edge_index - (2, E) ``[src, dst]`` edge endpoints (flat local indices). + (2, E) destination-major ``[src, dst]`` endpoints. edge_vec (E, 3) neighbor-minus-center edge vectors (sample for tracing). edge_mask (E,) valid-edge mask (sample for tracing). + destination_order + (E,) identity destination permutation. The exported ABI + requires this canonical layout. + source_order + (E,) source-grouped edge permutation. + destination_row_ptr, source_row_ptr + (N + 1,) destination/source CSR offsets. + destination_sorted + Static export-time assertion that the payload is + destination-major and ``destination_order`` is identity. fparam, aparam, do_atomic_virial, charge_spin As in ``forward_common_lower_graph``. **make_fx_kwargs @@ -704,18 +854,35 @@ def forward_common_lower_graph_exportable( ------- torch.nn.Module A traced module whose ``forward`` accepts - ``(atype, n_node, edge_index, edge_vec, edge_mask, - fparam, aparam, charge_spin)`` and returns a dict with the - same internal keys as ``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, fparam, aparam, charge_spin)`` and returns a + dict with the same internal keys as + ``forward_common_lower_graph``. """ + validate_graph_csr_for_export( + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + atype.shape[0], + destination_sorted=destination_sorted, + ) model = self def fn( atype: torch.Tensor, n_node: torch.Tensor, + n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, @@ -726,9 +893,15 @@ def fn( 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=destination_sorted, do_atomic_virial=do_atomic_virial, fparam=fparam, aparam=aparam, @@ -738,9 +911,14 @@ def fn( return make_fx(fn, **make_fx_kwargs)( 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, @@ -773,13 +951,10 @@ def forward_common_lower_exportable_with_comm( for GNN descriptors via the opaque ``deepmd_export::border_op`` wrapper. The comm tensors enter the exported program as 8 additional positional inputs after - the usual (coord, atype, nlist, mapping, fparam, aparam) — - this fixes the C++ ABI for ``DeepPotPTExpt`` (Phase 4). + the usual (coord, atype, nlist, mapping, fparam, aparam). - Tracing requires ``nswap >= 1`` (Phase 0 finding); with - ``nswap == 0`` the dim specializes and the artifact would - only run for that exact value. The C++ caller must always - provide ``nswap >= 1``. + Tracing requires ``nswap >= 1``; zero specializes the dimension. + The C++ caller must always provide ``nswap >= 1``. """ model = self diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f8dc2d9a1f..08218835d7 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -65,6 +65,9 @@ from deepmd.pt_expt.model import ( get_model, ) +from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower, +) from deepmd.pt_expt.train.wrapper import ( ModelWrapper, ) @@ -579,44 +582,6 @@ def _finalize_compiled_lower( ) -def _model_uses_graph_lower(model: torch.nn.Module) -> bool: - """Whether ``model``'s eager default-flip routes through the GRAPH lower. - - Mirrors the predicate in - :meth:`~deepmd.pt_expt.model.make_model.make_model..CM._resolve_graph_method` - for ``neighbor_graph_method is None`` (the training default): a model is - graph-eligible iff it is ``mixed_types`` AND its single descriptor reports - ``uses_graph_lower() == True`` (dpa1/se_atten with concat type embedding; - attention layers included). - - When True the compiled lower must be the GRAPH ``forward_common_lower_graph`` - so the compiled path matches eager training (which already default-flips to - the carry-all graph forward); when False the dense ``forward_lower`` is - compiled (se_e2_a / dpa2 / dpa3 / linear / zbl). - - ASSUMPTION: training uses the default ``neighbor_graph_method`` (None). If a - user-facing ``"legacy"`` opt-out is ever plumbed into the trainer, this gate - must also honor it (else eager would run dense while the compiled path runs - the graph lower, re-introducing the eager!=compiled divergence this fixes). - """ - if not hasattr(model, "mixed_types"): - return False - try: - if not model.mixed_types(): - return False - except (AttributeError, NotImplementedError): - return False - # Linear / ZBL atomic models have no single ``descriptor`` -> dense. - descriptor = getattr(getattr(model, "atomic_model", None), "descriptor", None) - uses_graph = getattr(descriptor, "uses_graph_lower", None) - if uses_graph is None: - return False - try: - return bool(uses_graph()) - except (AttributeError, NotImplementedError): - return False - - def _trace_and_compile_graph( model: torch.nn.Module, fparam: torch.Tensor | None, @@ -739,9 +704,14 @@ def _trace_and_compile_graph( ( s_atype, s_n_node, + s_n_local, s_edge_index, s_edge_vec, s_edge_mask, + s_destination_order, + s_destination_row_ptr, + s_source_order, + s_source_row_ptr, s_fparam, s_aparam, s_charge_spin, @@ -750,9 +720,14 @@ def _trace_and_compile_graph( def fn( atype: torch.Tensor, n_node: torch.Tensor, + n_local: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, @@ -778,9 +753,14 @@ def fn( model_ret = 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, do_atomic_virial=False, fparam=fparam, aparam=aparam, @@ -813,9 +793,14 @@ def fn( )( s_atype, s_n_node, + s_n_local, s_edge_index, s_edge_vec, s_edge_mask, + s_destination_order, + s_destination_row_ptr, + s_source_order, + s_source_row_ptr, s_fparam, s_aparam, s_charge_spin, @@ -913,7 +898,7 @@ def forward( # lower too, otherwise the eager (graph) and compiled (dense) backward # gradients diverge at fp64 accumulation and the optimizer amplifies it. if self._graph_eligible is None: - self._graph_eligible = _model_uses_graph_lower(self.original_model) + self._graph_eligible = model_uses_graph_lower(self.original_model) if self._graph_eligible: return self._forward_graph( coord, atype, box, fparam, aparam, charge_spin, nframes, nloc, rcut @@ -1222,9 +1207,14 @@ def _forward_graph( result = self.compiled_forward_lower( atype_flat, ng.n_node, + ng.n_node, ng.edge_index, edge_vec, ng.edge_mask, + ng.destination_order, + ng.destination_row_ptr, + ng.source_order, + ng.source_row_ptr, fparam, aparam, charge_spin, diff --git a/deepmd/pt_expt/utils/canonical_graph.py b/deepmd/pt_expt/utils/canonical_graph.py new file mode 100644 index 0000000000..5db225ab1a --- /dev/null +++ b/deepmd/pt_expt/utils/canonical_graph.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compact canonical graph contract for compressed DPA1 deployment.""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import torch + +if TYPE_CHECKING: + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + + +@dataclass(frozen=True) +class DPA1CanonicalGraph: + """Store a cutoff-compacted destination-major graph without redundant fields. + + The physical edge count is the final value of either CSR row-pointer tensor. + The storage edge axis has length ``max(E, 2)``; suffix entries are guards + excluded from every CSR row. + + Parameters + ---------- + n_node + Per-frame total node counts with shape ``(nf,)``, int64. + n_local + Per-frame owned node counts with shape ``(nf,)``, int64. + source + Source-node index for each edge storage slot with shape ``(S,)``, + int64. + edge_vec + Neighbor-minus-center vectors with shape ``(S, 3)``, float32. + destination_row_ptr + Destination CSR offsets with shape ``(N + 1,)``, int64. + source_row_ptr + Source CSR offsets with shape ``(N + 1,)``, int64. + source_order + Edge storage positions grouped by source with shape ``(S,)`` and the + same dtype as ``source``. + """ + + 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 + + +def validate_canonical_graph_shapes( + graph: DPA1CanonicalGraph, + 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: + 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") + if graph.n_node.dtype != torch.int64 or graph.n_local.dtype != torch.int64: + raise ValueError("canonical graph node counts must be int64") + if ( + graph.destination_row_ptr.dtype != torch.int64 + or graph.source_row_ptr.dtype != torch.int64 + ): + raise ValueError("canonical graph row pointers must be int64") + if graph.source.ndim != 1 or graph.source_order.shape != graph.source.shape: + raise ValueError("canonical graph source arrays must share shape (S,)") + if graph.edge_vec.shape != (graph.source.shape[0], 3): + 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.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,): + raise ValueError("source_row_ptr must have shape (N + 1,)") + if graph.n_node.shape != graph.n_local.shape or graph.n_node.ndim != 1: + raise ValueError("n_node and n_local must share shape (nf,)") + devices = { + graph.n_node.device, + graph.n_local.device, + graph.source.device, + graph.edge_vec.device, + graph.destination_row_ptr.device, + graph.source_row_ptr.device, + graph.source_order.device, + } + if len(devices) != 1: + raise ValueError("canonical graph tensors must reside on one device") + tensors = ( + graph.n_node, + graph.n_local, + graph.source, + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + ) + if not all(tensor.is_contiguous() for tensor in tensors): + raise ValueError("canonical graph tensors must be contiguous") + + +def canonical_graph_from_neighbor_graph( + graph: NeighborGraph, +) -> DPA1CanonicalGraph: + """Convert a compact generic graph into the source-only deployment contract. + + Parameters + ---------- + graph + Destination-major graph whose physical edges form a valid prefix. + + Returns + ------- + DPA1CanonicalGraph + Source-only graph with exactly two storage slots when ``E < 2``. + + Raises + ------ + ValueError + If the graph is not compact, canonical, or dual-CSR. + """ + if not graph.destination_sorted: + raise ValueError("canonical deployment requires destination-major payload") + if ( + graph.destination_row_ptr is None + or graph.source_order is None + or graph.source_row_ptr is None + ): + raise ValueError("canonical deployment requires destination/source CSR") + if graph.n_local is None: + raise ValueError("canonical deployment requires explicit owned-node counts") + + physical_edge_count = int(graph.destination_row_ptr[-1].item()) + if int(graph.source_row_ptr[-1].item()) != physical_edge_count: + raise ValueError("destination/source CSR physical edge counts differ") + if physical_edge_count > graph.edge_index.shape[1]: + raise ValueError("CSR physical edge count exceeds edge storage") + if not bool(torch.all(graph.edge_mask[:physical_edge_count]).item()): + raise ValueError("canonical deployment requires all physical edges active") + if bool(torch.any(graph.edge_mask[physical_edge_count:]).item()): + raise ValueError("canonical deployment guards must form an inactive suffix") + + node_count = graph.destination_row_ptr.shape[0] - 1 + + storage_edge_count = max(physical_edge_count, 2) + source = torch.zeros( + storage_edge_count, + dtype=torch.int64, + device=graph.edge_index.device, + ) + edge_vec = torch.zeros( + storage_edge_count, + 3, + dtype=torch.float32, + device=graph.edge_vec.device, + ) + source_order = torch.arange( + storage_edge_count, + dtype=torch.int64, + device=graph.edge_index.device, + ) + if physical_edge_count: + source[:physical_edge_count] = graph.edge_index[0, :physical_edge_count].to( + torch.int64 + ) + 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) + + result = DPA1CanonicalGraph( + n_node=graph.n_node.contiguous(), + n_local=graph.n_local.contiguous(), + source=source, + edge_vec=edge_vec, + destination_row_ptr=graph.destination_row_ptr.contiguous(), + source_row_ptr=graph.source_row_ptr.contiguous(), + source_order=source_order, + ) + validate_canonical_graph_shapes(result, node_count) + return result + + +__all__ = [ + "DPA1CanonicalGraph", + "canonical_graph_from_neighbor_graph", + "validate_canonical_graph_shapes", +] diff --git a/deepmd/pt_expt/utils/edge_schema.py b/deepmd/pt_expt/utils/edge_schema.py index 80871fdee9..baa7eca2dd 100644 --- a/deepmd/pt_expt/utils/edge_schema.py +++ b/deepmd/pt_expt/utils/edge_schema.py @@ -257,10 +257,15 @@ def edge_schema_from_ij_shifts( torch.any(shifts != 0, dim=1), as_tuple=False ).flatten() if shifted_idx.numel() > 0: + # Image offset ``S @ cell`` as a broadcast multiply-reduce: the + # (n_shift, 3) @ (3, 3) matmul otherwise dispatches to an fp64 GEMM + # kernel whose length-3 contraction is catastrophically inefficient, + # while this stays bit-identical. + sel_shifts = shifts.index_select(0, shifted_idx) edge_vec_all.index_add_( 0, shifted_idx, - shifts.index_select(0, shifted_idx) @ cell, + (sel_shifts[:, :, None] * cell).sum(1), ) edge_len2 = torch.sum(edge_vec_all * edge_vec_all, dim=-1) edge_keep = (edge_len2 > 1e-10) & (edge_len2 <= float(rcut) * float(rcut)) diff --git a/deepmd/pt_expt/utils/graph_csr.py b/deepmd/pt_expt/utils/graph_csr.py new file mode 100644 index 0000000000..1a0484931a --- /dev/null +++ b/deepmd/pt_expt/utils/graph_csr.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""PyTorch validation for graph-lower CSR export inputs.""" + +from __future__ import ( + annotations, +) + +import torch + + +def _validate_row_ptr( + name: str, + row_ptr: torch.Tensor, + node_count: int, + edge_count: int, + device: torch.device, +) -> None: + """Validate the bounds and layout of one CSR row-pointer tensor.""" + if row_ptr.dtype != torch.int64 or row_ptr.shape != (node_count + 1,): + raise ValueError( + f"graph export requires {name} with shape (N + 1,) and dtype int64" + ) + if ( + row_ptr.device != device + or not bool(torch.all(row_ptr[1:] >= row_ptr[:-1])) + or int(row_ptr[0]) != 0 + or int(row_ptr[-1]) > edge_count + ): + raise ValueError(f"graph export received invalid {name}") + + +def _validate_permutation_csr( + name: str, + order: torch.Tensor, + row_ptr: torch.Tensor, + endpoint: torch.Tensor, + edge_mask: torch.Tensor, + node_count: int, + edge_count: int, + device: torch.device, +) -> None: + """Validate that a CSR order covers every active edge in its endpoint row.""" + if order.dtype not in (torch.int32, torch.int64) or order.shape != (edge_count,): + raise ValueError( + f"graph export requires {name} with shape (E,) and integer dtype" + ) + if order.device != device: + raise ValueError(f"graph export requires {name} on the edge device") + expected = torch.arange(edge_count, dtype=order.dtype, device=device) + if not torch.equal(torch.sort(order).values, expected): + raise ValueError(f"graph export requires {name} to be an edge permutation") + + active_edge = torch.arange( + edge_count, + dtype=torch.int64, + device=device, + )[edge_mask] + active_endpoint = endpoint[active_edge] + if active_endpoint.numel() and ( + int(active_endpoint.min()) < 0 or int(active_endpoint.max()) >= node_count + ): + raise ValueError("graph export received an active edge outside node bounds") + + inverse = torch.empty(edge_count, dtype=torch.int64, device=device) + inverse.scatter_( + 0, + order.to(torch.int64), + torch.arange(edge_count, dtype=torch.int64, device=device), + ) + active_position = inverse[active_edge] + row_begin = row_ptr[active_endpoint] + row_end = row_ptr[active_endpoint + 1] + if not bool( + torch.all((row_begin <= active_position) & (active_position < row_end)) + ): + raise ValueError(f"graph export received {name} entries outside their CSR rows") + + +def validate_graph_csr_for_export( + 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, + node_count: int, + *, + destination_sorted: bool, +) -> None: + """Validate concrete graph CSR inputs before graph-form export tracing. + + Parameters + ---------- + edge_index : torch.Tensor + Edge endpoints with shape ``(2, E)`` and integer dtype. + edge_mask : torch.Tensor + Valid-edge mask with shape ``(E,)`` and dtype bool. + destination_order, source_order : torch.Tensor + Edge permutations with shape ``(E,)``. + destination_row_ptr, source_row_ptr : torch.Tensor + CSR offsets with shape ``(N + 1,)`` and dtype int64. + node_count : int + Number of nodes on the flat graph axis. + destination_sorted : bool + Whether the exported descriptor addresses destination rows directly. + + Raises + ------ + ValueError + If an order is not a permutation, a CSR row omits an active edge, or a + canonical export does not use the identity destination order. + + Notes + ----- + This function validates concrete trace inputs only. It must execute before + ``make_fx`` or ``torch.export`` so it does not add inference operations. + """ + if edge_index.ndim != 2 or edge_index.shape[0] != 2: + raise ValueError("graph export requires edge_index with shape (2, E)") + if edge_index.dtype not in (torch.int32, torch.int64): + raise ValueError("graph export requires integer edge_index") + edge_count = edge_index.shape[1] + if ( + edge_mask.shape != (edge_count,) + or edge_mask.dtype != torch.bool + or edge_mask.device != edge_index.device + ): + raise ValueError("graph export requires a bool edge_mask with shape (E,)") + + _validate_row_ptr( + "destination_row_ptr", + destination_row_ptr, + node_count, + edge_count, + edge_index.device, + ) + _validate_row_ptr( + "source_row_ptr", + source_row_ptr, + node_count, + edge_count, + edge_index.device, + ) + _validate_permutation_csr( + "destination_order", + destination_order, + destination_row_ptr, + edge_index[1], + edge_mask, + node_count, + edge_count, + edge_index.device, + ) + _validate_permutation_csr( + "source_order", + source_order, + source_row_ptr, + edge_index[0], + edge_mask, + node_count, + edge_count, + edge_index.device, + ) + if destination_sorted: + expected = torch.arange( + edge_count, + dtype=destination_order.dtype, + device=destination_order.device, + ) + if not torch.equal(destination_order, expected): + raise ValueError( + "destination_sorted=True requires an identity destination_order" + ) diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py index fe588c94e7..637dde9912 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -37,6 +37,7 @@ GraphLayout, NeighborGraph, apply_pair_exclusion, + attach_edge_csr, neighbor_graph_from_ijs, ) from deepmd.pt.utils.nv_nlist import ( @@ -212,6 +213,8 @@ def build_neighbor_graph_nv( rcut: float, layout: GraphLayout | None = None, *, + with_csr: bool = False, + canonicalize: bool = False, pair_excl: PairExcludeMask | None = None, compact: bool = False, ) -> NeighborGraph: @@ -229,6 +232,12 @@ def build_neighbor_graph_nv( cutoff radius. layout edge-axis length policy; ``None`` => dynamic with ``min_edges`` guards. + with_csr + Whether to construct destination/source CSR views for a consumer that + requires edge-grouped reductions. + canonicalize + Whether to reorder every edge field into destination-major form. Implies + ``with_csr=True``. pair_excl Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` for model-level ``pair_exclude_types``. When given, @@ -261,6 +270,9 @@ def build_neighbor_graph_nv( "install with `pip install nvalchemi-toolkit-ops` or use " "neighbor_graph_method='dense'." ) + from nvalchemiops.neighbors.neighbor_utils import ( + estimate_max_neighbors, + ) device = coord.device nf = coord.shape[0] if coord.ndim == 3 else 1 @@ -272,7 +284,16 @@ def build_neighbor_graph_nv( empty_i = torch.zeros((0,), dtype=torch.int64, device=device) empty_S = torch.zeros((0, 3), dtype=torch.int64, device=device) return neighbor_graph_from_ijs( - empty_i, empty_i, empty_S, coord, box, empty_i, nloc, layout=layout + empty_i, + empty_i, + empty_S, + coord, + box, + empty_i, + nloc, + layout=layout, + with_csr=with_csr, + canonicalize=canonicalize, ) # Carry-all: grow capacity until every neighbor fits (no sel cap). @@ -280,8 +301,12 @@ def build_neighbor_graph_nv( # vesin handles unwrapped positions natively), nvalchemiops requires # in-cell positions, so BOTH the search and the edge_vec recomputation use # the normalized coords; S then matches the coords the search actually saw. + initial_capacity = max( + 64, + estimate_max_neighbors(float(rcut), safety_factor=1.25), + ) coord, cell, neighbor_matrix, num_neighbors, shifts = nv_search_matrix( - coord, box, rcut, start_capacity=max(64, nloc) + coord, box, rcut, start_capacity=initial_capacity ) box_out = cell # edge_vec is recomputed from these (normalized) coords @@ -300,9 +325,18 @@ def build_neighbor_graph_nv( shift, frame_idx = shift[keep], frame_idx[keep] graph = neighbor_graph_from_ijs( - center_local, src_local, shift, coord, box_out, frame_idx, nloc, layout=layout + center_local, + src_local, + shift, + coord, + box_out, + frame_idx, + nloc, + layout=layout, ) if pair_excl is not None: at_flat = torch.as_tensor(atype, device=device).reshape(-1) graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact) + if with_csr or canonicalize: + graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize) return graph diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index c6d26dcde8..fe8357b858 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1,6 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import contextlib import ctypes import json +import logging +import os +from collections.abc import ( + Iterator, +) from typing import ( Any, ) @@ -12,6 +18,8 @@ build_neighbor_list, extend_coord_with_ghosts, ) + +log = logging.getLogger(__name__) from deepmd.dpmodel.utils.region import ( normalize_coord, ) @@ -50,22 +58,19 @@ def _strip_shape_assertions(graph_module: torch.nn.Module) -> None: (erasing can disturb the FX graph and yield NaN gradients on some torch versions). - Called from TWO export paths in ``_trace_and_export``: + Called from two export paths in ``_trace_and_export``: * **spin (dense) models** — atom-doubling slice patterns depend on ``(nall - nloc)``, producing spurious guards like ``Ne(nall, nloc)``; the model is correct even when ``nall == nloc`` (NoPBC, no ghosts). - * **graph models** — the DYNAMIC edge axis (``Dim("nedge")``) produces guards - of the ``nloc_min``/SIGFPE family on the edge count ``E``. These are the - shape-specialization guards the static-``edge_capacity`` path was designed - to avoid; neutralising them is what makes one artifact eval any edge count. + * **graph models** — the dynamic edge axis (``Dim("nedge")``) produces + shape-specialization guards on the edge count ``E``. - **Safety:** in both contexts every input is constructed well-formed by the + In both contexts every input is constructed well-formed by the builder (spin: valid atom doubling; graph: ``build_neighbor_graph`` / ``buildGraphTensors`` always emit ``E >= min_edges == 2`` with in-range, - masked edges), so the neutralised guards would never legitimately fire. The - only cost is that a MALFORMED runtime tensor no longer throws cleanly — the - documented AOTI trade-off (CLAUDE.md), accepted identically on both paths. + masked edges). Malformed runtime tensors are outside this exported ABI and + are not guaranteed to trigger these shape assertions. """ graph = graph_module.graph for node in list(graph.nodes): @@ -203,8 +208,8 @@ def _make_comm_sample_inputs( ) -> tuple[torch.Tensor, ...]: """Build trivial-but-valid comm tensors for tracing the with-comm variant. - Phase 0 finding: tracing with ``nswap == 0`` causes the dim to - specialize, so we must use ``nswap >= 1``. We use ``nswap == 1`` + Tracing with ``nswap == 0`` specializes the dimension, so the sample uses + ``nswap == 1`` with a single self-send swap whose sendlist points to ``nghost`` local atoms (the actual indices don't matter for the trace — only the validity of the pointer matters; ``border_op`` is opaque to @@ -367,6 +372,7 @@ def build_synthetic_graph_inputs( nloc: int = 7, *, dtype: torch.dtype, + edge_dtype: torch.dtype | None = None, device: torch.device | None = None, want_fparam: bool = True, want_aparam: bool = True, @@ -380,10 +386,13 @@ def build_synthetic_graph_inputs( traces can never desync on the graph input schema. Builds a small random system, runs the carry-all :func:`~deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph` with a - STATIC ``GraphLayout(edge_capacity=e_max)`` (decision #16: the masked static - edge axis), and returns tensors in the positional order expected by + padded ``GraphLayout(edge_capacity=e_max)`` trace sample, then canonicalizes + it to the destination-major deployment ABI. The exported edge axis remains + dynamic; the concrete capacity only supplies representative tensors to + ``make_fx``. Inputs follow the positional order expected by ``forward_(common_)lower_graph``: - ``(atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, + ``(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)``. The system (``rng(42)``, ``box = rcut*3``, centered coords, ``atype[:, i] = @@ -395,15 +404,18 @@ def build_synthetic_graph_inputs( model : torch.nn.Module The pt_expt energy model (must expose ``get_rcut``/``get_type_map``/...). e_max : int - Static edge capacity ``E`` to pad the (masked) edge axis to. + Concrete edge-axis size used by the trace sample. nframes : int Number of frames in the sample system. nloc : int Number of local atoms per frame (``N == nframes * nloc``). dtype : torch.dtype - Float precision of ``coord``/``edge_vec``/``fparam``/... . The exported - ``.pt2`` is float64-only (C++ ABI); training passes - ``GLOBAL_PT_FLOAT_PRECISION``. + Float precision of ``coord``/``fparam`` and other conditioning inputs. + edge_dtype : torch.dtype, optional + Precision of ``edge_vec``. Defaults to ``dtype``. A compressed DPA1 + graph artifact uses float32 because its descriptor and analytical + backward both compute in float32; generic graph artifacts preserve the + model input precision. device : torch.device, optional Target device. Defaults to ``deepmd.pt_expt.utils.env.DEVICE``; the export path passes ``cpu`` explicitly (make_fx traces on CPU). @@ -420,6 +432,8 @@ def build_synthetic_graph_inputs( if device is None: device = _env.DEVICE + if edge_dtype is None: + edge_dtype = dtype rcut = model.get_rcut() ntypes = len(model.get_type_map()) @@ -440,7 +454,12 @@ def build_synthetic_graph_inputs( atype_t = torch.tensor(atype_np, dtype=torch.int64, device=device) box_t = torch.tensor(np.tile(box_np, (nframes, 1)), dtype=dtype, device=device) graph = build_neighbor_graph( - coord_t, atype_t, box_t, rcut, layout=GraphLayout(edge_capacity=e_max) + coord_t, + atype_t, + box_t, + rcut, + layout=GraphLayout(edge_capacity=e_max), + canonicalize=True, ) fparam = ( @@ -458,41 +477,135 @@ def build_synthetic_graph_inputs( if (want_charge_spin and dim_chg_spin > 0) else None ) + # Keep total and owned counts value-distinct during tracing so export does + # not specialize the multi-rank ownership relation to ``n_local == n_node``. + n_local = torch.clamp(graph.n_node - 1, min=1) return ( atype_t.reshape(-1), graph.n_node, + n_local, graph.edge_index, - graph.edge_vec, + graph.edge_vec.to(edge_dtype), graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, fparam, aparam, charge_spin, ) +def build_synthetic_canonical_graph_inputs( + model: torch.nn.Module, + e_max: int, + *, + device: torch.device, +) -> tuple[torch.Tensor, ...]: + """Build the compact canonical trace inputs for compressed DPA1.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + ) + + sample = build_synthetic_graph_inputs( + model, + e_max, + dtype=torch.float32, + edge_dtype=torch.float32, + device=device, + want_fparam=False, + want_aparam=False, + want_charge_spin=False, + ) + ( + 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 + graph = NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + n_local=n_local, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=True, + ) + compact = canonical_graph_from_neighbor_graph(graph) + return ( + atype, + compact.n_node, + compact.n_local, + compact.source, + compact.edge_vec, + compact.destination_row_ptr, + compact.source_row_ptr, + compact.source_order, + ) + + +def _build_canonical_graph_dynamic_shapes( + *sample_inputs: torch.Tensor, +) -> tuple: + """Build dynamic shapes for the eight-tensor compact deployment ABI.""" + del sample_inputs + 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 ( + {0: node_dim}, + {0: nframes_dim}, + {0: nframes_dim}, + {0: edge_storage_dim}, + {0: edge_storage_dim}, + {0: node_dim + 1}, + {0: node_dim + 1}, + {0: edge_storage_dim}, + ) + + def _build_graph_dynamic_shapes( *sample_inputs: torch.Tensor | None, ) -> tuple: """Build dynamic-shape specifications for the graph-form forward_lower export. - ``nframes`` (the ``n_node`` axis), ``N`` (the flat node axis) AND the edge - axis ``E`` are all dynamic dims (B2.0: the dynamic edge axis replaces the - static ``edge_capacity`` of B1). ``E`` is marked ``Dim("nedge", min=2)`` so - the AOTI artifact accepts any system size with no capacity ceiling — the - ``min=2`` lower bound mirrors the dense path's ``Dim("nnei", min=...)`` (a - dynamic, SIGFPE-tamed axis) and matches the carry-all builder's + ``nframes`` (the ``n_node`` axis), ``N`` (the flat node axis), and the edge + axis ``E`` are all dynamic dimensions. ``E`` is marked + ``Dim("nedge", min=2)`` so + the AOTI artifact accepts any system size with no capacity ceiling. The + ``min=2`` lower bound mirrors the dense path's ``Dim("nnei", min=...)`` and + matches the carry-all builder's ``min_edges=2`` guard (every dynamic graph carries >=2 edges). Parameters ---------- *sample_inputs : torch.Tensor | None - ``(atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, - charge_spin)`` — 8 entries matching ``forward_lower_graph_exportable``. + ``(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)`` — 13 entries matching + ``forward_lower_graph_exportable``. """ - fparam = sample_inputs[5] - aparam = sample_inputs[6] - charge_spin = sample_inputs[7] + fparam = sample_inputs[10] + aparam = sample_inputs[11] + charge_spin = sample_inputs[12] nframes_dim = torch.export.Dim("nframes", min=1) n_node_total_dim = torch.export.Dim("n_node_total", min=1) nedge_dim = torch.export.Dim("nedge", min=2) @@ -500,9 +613,14 @@ def _build_graph_dynamic_shapes( return ( {0: n_node_total_dim}, # atype: (N,) {0: nframes_dim}, # n_node: (nf,) + {0: nframes_dim}, # n_local: (nf,) {1: nedge_dim}, # edge_index: (2, E) — E dynamic {0: nedge_dim}, # edge_vec: (E, 3) — E dynamic {0: nedge_dim}, # edge_mask: (E,) — E dynamic + {0: nedge_dim}, # destination_order: (E,) + {0: n_node_total_dim + 1}, # destination_row_ptr: (N + 1,) + {0: nedge_dim}, # source_order: (E,) + {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) {0: nframes_dim} if fparam is not None else None, # fparam: (nf, ndf) # aparam: (nf, nloc, nda) — both the frame AND atom axes are dynamic, # matching the dense ``_build_dynamic_shapes`` (otherwise a dim_aparam>0 @@ -618,8 +736,46 @@ def _build_dynamic_shapes( return (*base, None, None, None, None, None, None, None, None) +def _graph_edge_dtype(model: torch.nn.Module, lower_kind: str) -> str: + """Return the graph edge-vector dtype encoded by the deployment artifact. + + Geometrically compressed DPA1 with float32 descriptor statistics evaluates + both descriptor directions in float32 and therefore accepts float32 + geometry directly. Other graph descriptors retain the model-agnostic + float64 geometry ABI. + """ + atomic_model = getattr(model, "atomic_model", None) + descriptor = getattr(atomic_model, "descriptor", None) + descriptor_block = getattr(descriptor, "se_atten", None) + statistics = getattr(descriptor_block, "mean", None) + if ( + lower_kind in ("graph", "dpa1_canonical") + and bool(getattr(descriptor, "geo_compress", False)) + and isinstance(statistics, torch.Tensor) + and statistics.dtype == torch.float32 + ): + return "float32" + return "float64" + + +def _supports_graph_export(model: torch.nn.Module) -> bool: + """Whether the model has an exportable graph-lower implementation. + + A compressed descriptor must use its opaque graph operator during export; + tracing through the reference tabulation kernel is unsupported. + """ + atomic_model = getattr(model, "atomic_model", None) + descriptor = getattr(atomic_model, "descriptor", None) + if not bool(getattr(descriptor, "geo_compress", False)): + return True + eligible = getattr(descriptor, "_fused_eligible", None) + return callable(eligible) and bool(eligible("cuda")) + + def _collect_metadata( - model: torch.nn.Module, is_spin: bool = False, lower_kind: str = "nlist" + model: torch.nn.Module, + is_spin: bool = False, + lower_kind: str = "nlist", ) -> dict: """Collect metadata from the model for C++ inference. @@ -737,7 +893,8 @@ def _probe_has_message_passing(obj: object) -> bool | None: # "nlist" → dense quartet (extended_coord, extended_atype, nlist, mapping) # "graph" → NeighborGraph (atype, n_node, edge_index, edge_vec, edge_mask) # The C++ loader branches on this to build the matching inputs. - meta["lower_input_kind"] = "graph" if lower_kind == "graph" else "nlist" + meta["lower_input_kind"] = lower_kind + meta["graph_edge_dtype"] = _graph_edge_dtype(model, lower_kind) # Model-level pair-type exclusion (``pair_exclude_types``): a list of # ``[ti, tj]`` type pairs whose interaction is dropped. Exclusion is a @@ -825,6 +982,63 @@ def _serialize_from_file_pt2(model_file: str) -> dict: return model_dict +@contextlib.contextmanager +def _cuda_infer_at_least_2() -> Iterator[None]: + """Pin ``DP_CUDA_INFER`` to at least 2 for the duration of a trace. + + Level 2 emits the inference pipeline as explicit descriptor, fitting, + descriptor-backward, and CSR force/virial custom operators. These operators + remain opaque through ``torch.export``. The level-1 autograd lower can + decompose the analytic backward to aten, while level 0 selects the + untraceable reference tabulation. Level 2 degrades internally when an + operator is unavailable or ineligible, so it is a safe floor for graph + export. + """ + from deepmd.kernels.utils import ( + cuda_infer_level, + ) + + saved = os.environ.get("DP_CUDA_INFER") + if cuda_infer_level() < 2: + os.environ["DP_CUDA_INFER"] = "2" + try: + yield + finally: + if saved is None: + os.environ.pop("DP_CUDA_INFER", None) + else: + os.environ["DP_CUDA_INFER"] = saved + + +def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str: + """Resolve ``lower_kind="auto"`` to a concrete lower-forward schema. + + ``"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 + unchanged. + """ + if lower_kind != "auto": + return lower_kind + if not model_file.endswith(".pt2") or data["model"].get("type") == "spin_ener": + return "nlist" + from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower, + ) + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + 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, + ) + + return "dpa1_canonical" if canonical_model_eligible(model) else "graph" + return "nlist" + + def deserialize_to_file( model_file: str, data: dict, @@ -853,25 +1067,50 @@ def deserialize_to_file( tracing the uncompressed model (make_fx cannot trace custom ops). do_atomic_virial : bool If True, export with per-atom virial correction (3 extra backward - passes, ~2.5x slower). Default False for best performance. + passes, ~2.5x slower). Default False for best performance. Forced True + for a graph lower, whose LAMMPS Kokkos consumer always reads it. lower_kind : str Which lower-forward schema the compiled AOTI graph consumes: ``"nlist"`` (default) traces the dense quartet (``extended_coord``/``extended_atype``/``nlist``/``mapping``); ``"graph"`` traces the NeighborGraph schema - (``atype``/``n_node``/``edge_index``/``edge_vec``/``edge_mask``) with a - DYNAMIC edge axis ``E`` (``Dim("nedge", min=2)``), so the artifact - accepts any system size. The selected schema is recorded as - ``lower_input_kind`` in ``metadata.json``. + (``atype``/``n_node``/``edge_index``/``edge_vec``/``edge_mask`` and + the destination/source CSR views) with a DYNAMIC edge axis ``E`` + (``Dim("nedge", min=2)``), so the artifact accepts any system size. + ``"auto"`` (used by ``convert-backend``) resolves to ``"graph"`` for an + exportable graph-lower ``.pt2`` and ``"nlist"`` otherwise (see + :func:`_resolve_lower_kind`). A graph lower always preserves the fused + inference operators (``DP_CUDA_INFER >= 2``) and the per-atom virial. + The selected schema is recorded as ``lower_input_kind`` in + ``metadata.json``. """ - if model_file.endswith(".pt2"): - _deserialize_to_file_pt2( - model_file, data, model_json_override, do_atomic_virial, lower_kind - ) + lower_kind = _resolve_lower_kind(model_file, data, lower_kind) + # 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"): + do_atomic_virial = True + ctx: contextlib.AbstractContextManager = _cuda_infer_at_least_2() else: - _deserialize_to_file_pte( - model_file, data, model_json_override, do_atomic_virial, lower_kind - ) + ctx = contextlib.nullcontext() + with ctx: + if model_file.endswith(".pt2"): + _deserialize_to_file_pt2( + model_file, + data, + model_json_override, + do_atomic_virial, + lower_kind, + ) + else: + _deserialize_to_file_pte( + model_file, + data, + model_json_override, + do_atomic_virial, + lower_kind, + ) def _trace_and_export( @@ -944,15 +1183,29 @@ def _trace_and_export( model = BaseModel.deserialize(data["model"]) model.to("cpu") model.eval() + if lower_kind == "graph" and not _supports_graph_export(model): + raise NotImplementedError( + "graph-form export of a compressed descriptor requires its " + "float32 fused graph operator; use lower_kind='nlist' for this model" + ) + + # Autotune checkpoint-specific custom-kernel launch tables on the target + # GPU before tracing. The model itself remains on CPU for tracing. + from deepmd.kernels.autotune import ( + run_autotune, + ) + + run_autotune(model, target_device) # 2. Collect metadata - metadata = _collect_metadata(model, is_spin=is_spin, lower_kind=lower_kind) + metadata = _collect_metadata( + model, + is_spin=is_spin, + lower_kind=lower_kind, + ) - # 2b. Graph-form export branch (NeighborGraph schema). The graph path is - # LOCAL-only (no ghosts), single-rank, energy-model only in PR-A/PR-B; it - # traces ``forward_lower_graph_exportable`` with a DYNAMIC edge axis (B2.0). - # The dense (nlist) path below is left byte-unchanged. - if lower_kind == "graph": + # Graph-form exports use a dynamic edge axis and an energy-model contract. + if lower_kind in ("graph", "dpa1_canonical"): import math check_graph_trace_torch_version(model) @@ -963,73 +1216,73 @@ def _trace_and_export( if with_comm_dict: raise NotImplementedError( "graph-form .pt2 export does not support the with-comm artifact " - "(multi-rank graph message passing is a later PR)" + "required for multi-rank message passing" ) - if not hasattr(model, "forward_lower_graph_exportable"): + canonical = lower_kind == "dpa1_canonical" + required_method = ( + "forward_lower_canonical_graph_exportable" + if canonical + else "forward_lower_graph_exportable" + ) + if not hasattr(model, required_method): raise NotImplementedError( - f"model {type(model).__name__} has no " - "forward_lower_graph_exportable; graph-form .pt2 export " - "requires an energy model" + f"model {type(model).__name__} has no {required_method}" + ) + if canonical: + from deepmd.kernels.cuda.dpa1.canonical import ( + canonical_model_eligible, ) - # The edge axis is DYNAMIC (B2.0): the AOTI artifact accepts any edge - # count, so there is no capacity to bake. The trace sample is built at a - # concrete, padded edge size only to keep the trace tensors distinct - # from the other dynamic dims (nframes=2, N=14) under torch.export's - # duck-sizing; the value itself does NOT constrain runtime. + if not canonical_model_eligible(model): + raise NotImplementedError( + "compact canonical export requires an eligible compressed " + "DPA1 energy model" + ) + nloc_sample = 7 nnei = sum(model.get_sel()) e_sample = math.ceil(1.25 * nloc_sample * nnei) - - # make_fx traces on CPU; the .pt2 C++ ABI is float64-only. Pass device - # and dtype explicitly instead of mutating the module-level env.DEVICE. - sample_inputs = build_synthetic_graph_inputs( - model, - e_max=e_sample, - nframes=2, - nloc=nloc_sample, - dtype=torch.float64, - device=torch.device("cpu"), - ) - - ( - atype_g, - n_node_g, - edge_index_g, - edge_vec_g, - edge_mask_g, - fparam_g, - aparam_g, - charge_spin_g, - ) = sample_inputs - - # Trace via make_fx on CPU (decomposes autograd.grad into aten ops). - traced = model.forward_lower_graph_exportable( - atype_g, - n_node_g, - edge_index_g, - edge_vec_g, - edge_mask_g, - fparam=fparam_g, - aparam=aparam_g, - do_atomic_virial=do_atomic_virial, - charge_spin=charge_spin_g, - tracing_mode="symbolic", - _allow_non_fake_inputs=True, - ) - sample_out = traced( - atype_g, - n_node_g, - edge_index_g, - edge_vec_g, - edge_mask_g, - fparam_g, - aparam_g, - charge_spin_g, - ) + if canonical: + sample_inputs = build_synthetic_canonical_graph_inputs( + model, + e_sample, + device=torch.device("cpu"), + ) + traced = model.forward_lower_canonical_graph_exportable( + *sample_inputs, + do_atomic_virial=do_atomic_virial, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + dynamic_shapes = _build_canonical_graph_dynamic_shapes(*sample_inputs) + else: + edge_dtype = ( + torch.float32 + if metadata["graph_edge_dtype"] == "float32" + else torch.float64 + ) + sample_inputs = build_synthetic_graph_inputs( + model, + e_max=e_sample, + nframes=2, + nloc=nloc_sample, + dtype=torch.float64, + edge_dtype=edge_dtype, + device=torch.device("cpu"), + ) + traced = model.forward_lower_graph_exportable( + *sample_inputs[:10], + fparam=sample_inputs[10], + aparam=sample_inputs[11], + do_atomic_virial=do_atomic_virial, + charge_spin=sample_inputs[12], + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + dynamic_shapes = _build_graph_dynamic_shapes(*sample_inputs) + sample_out = traced(*sample_inputs) output_keys = list(sample_out.keys()) - - dynamic_shapes = _build_graph_dynamic_shapes(*sample_inputs) exported = torch.export.export( traced, sample_inputs, @@ -1041,11 +1294,9 @@ def _trace_and_export( # Neutralise shape-guard assertion nodes on the dynamic edge axis. # ``prefer_deferred_runtime_asserts_over_guards=True`` converts the # symbolic-shape guards discovered while tracing into deferred - # ``aten._assert_scalar`` nodes; on the dynamic ``E`` axis these are the - # SIGFPE-prone ``nloc_min``-family checks (CLAUDE.md AOTI pitfalls) that - # the dense spin path already strips. Replacing each condition with - # ``True`` (not erasing the node) keeps the graph well-formed while - # letting the AOTI artifact generalise across edge counts. + # ``aten._assert_scalar`` nodes. Replacing each condition with ``True`` + # preserves graph structure while allowing the AOTI artifact to + # generalise across edge counts. _strip_shape_assertions(exported.graph_module) if target_device.type != "cpu": @@ -1056,7 +1307,7 @@ def _trace_and_export( exported = move_to_device_pass(exported, target_device) metadata["do_atomic_virial"] = do_atomic_virial - # The edge axis is DYNAMIC (B2.0): the AOTI forward accepts any edge + # The AOTI forward accepts any edge # count, so there is no ``edge_capacity`` to persist. The C++ / Python # conversion hub builds the carry-all graph at its exact (tight) edge # count and feeds it straight through. @@ -1164,6 +1415,7 @@ def _trace_and_export( # 4. Trace via make_fx on CPU. # This decomposes torch.autograd.grad into aten ops so the resulting # GraphModule no longer contains autograd calls. + log.info("Tracing the lower graph on CPU (make_fx)...") if is_spin: if with_comm_dict: traced = model.forward_common_lower_exportable_with_comm( @@ -1234,6 +1486,7 @@ def _trace_and_export( # graph. Exporting on CPU keeps devices consistent; we move the # ExportedProgram to the target device afterwards via the official # move_to_device_pass (avoids FakeTensor device-propagation errors). + log.info("Exporting the traced graph (torch.export)...") dynamic_shapes = _build_dynamic_shapes( *sample_inputs, has_spin=is_spin, @@ -1374,21 +1627,36 @@ def _deserialize_to_file_pt2( # compilations from multiple threads would race on this global. Callers # must serialise ``.pt2`` exports if running under a thread pool. # Processes are fine (each has its own inductor config). - import torch._inductor.config as _inductor_config - import deepmd.pt_expt.utils.env as _env + from deepmd.pt.utils.compile_compat import ( + build_inductor_compile_options, + patch_inductor_force_int64_indexing, + ) is_cuda = _env.DEVICE.type == "cuda" - saved_threshold = _inductor_config.realize_opcount_threshold - saved_assert_indexing = _inductor_config.assert_indirect_indexing + # Force int64 tensor indexing so the flattened index of a large + # data-dependent tensor never wraps past 2**31 into an illegal address. + patch_inductor_force_int64_indexing() + # The AOTInductor freeze must use the same Inductor lockdown as the + # pt-backend compile -- most importantly ``triton.max_tiles = 1``, which + # keeps the data-dependent edge / node axis on the x launch dimension + # (limit 2**31-1) rather than a 2-D y/z tile (limit 65535). Without it a + # compressed graph .pt2 (its level-1 lower carries an aten glue over the + # ``(n_node, NG * axis)`` descriptor) launches an out-of-range grid and + # fails at runtime with a CUDA "invalid argument" once that tensor exceeds + # 2**22 elements. The two deepmd-specific relaxations layer on top. + aoti_configs = build_inductor_compile_options(inference=True) + aoti_configs["assert_indirect_indexing"] = False if is_cuda: - _inductor_config.realize_opcount_threshold = 0 - _inductor_config.assert_indirect_indexing = False - try: - aoti_compile_and_package(exported, package_path=model_file) - finally: - _inductor_config.realize_opcount_threshold = saved_threshold - _inductor_config.assert_indirect_indexing = saved_assert_indexing + aoti_configs["realize_opcount_threshold"] = 0 + log.info( + "Compiling the AOTInductor package for %s (the slowest freeze stage; " + "typically several minutes)...", + _env.DEVICE, + ) + aoti_compile_and_package( + exported, package_path=model_file, inductor_configs=aoti_configs + ) # Second artifact: with-comm. Only for descriptors whose message # passing extends across rank boundaries. The flag was computed @@ -1406,16 +1674,9 @@ def _deserialize_to_file_pt2( ) with tempfile.TemporaryDirectory() as td: wc_path = os.path.join(td, "forward_lower_with_comm.pt2") - saved_threshold = _inductor_config.realize_opcount_threshold - saved_assert_indexing = _inductor_config.assert_indirect_indexing - if is_cuda: - _inductor_config.realize_opcount_threshold = 0 - _inductor_config.assert_indirect_indexing = False - try: - aoti_compile_and_package(exported_wc, package_path=wc_path) - finally: - _inductor_config.realize_opcount_threshold = saved_threshold - _inductor_config.assert_indirect_indexing = saved_assert_indexing + aoti_compile_and_package( + exported_wc, package_path=wc_path, inductor_configs=aoti_configs + ) with open(wc_path, "rb") as f: with_comm_bytes = f.read() # The output keys are identical between the two artifacts (same diff --git a/deepmd/pt_expt/utils/vesin_graph_builder.py b/deepmd/pt_expt/utils/vesin_graph_builder.py index 414b356f68..a715189ae5 100644 --- a/deepmd/pt_expt/utils/vesin_graph_builder.py +++ b/deepmd/pt_expt/utils/vesin_graph_builder.py @@ -35,6 +35,7 @@ GraphLayout, NeighborGraph, apply_pair_exclusion, + attach_edge_csr, neighbor_graph_from_ijs, ) from deepmd.pt_expt.utils.vesin_neighbor_list import ( @@ -97,6 +98,8 @@ def build_neighbor_graph_vesin( rcut: float, layout: GraphLayout | None = None, *, + with_csr: bool = False, + canonicalize: bool = False, pair_excl: PairExcludeMask | None = None, compact: bool = False, ) -> NeighborGraph: @@ -107,16 +110,21 @@ def build_neighbor_graph_vesin( Parameters ---------- - coord - (nf, nloc, 3) or (nf, nloc*3) local coordinates (torch tensor). - atype - (nf, nloc) local atom types; ``type < 0`` marks a virtual atom. - box - (nf, 3, 3) simulation cell, or ``None`` for non-periodic. - rcut - cutoff radius. - layout - edge-axis length policy; ``None`` => dynamic with ``min_edges`` guards. + coord : torch.Tensor + Coordinates with shape ``(nf, nloc, 3)``. + atype : torch.Tensor + Atom types with shape ``(nf, nloc)``. + box : torch.Tensor or None + Simulation cells with shape ``(nf, 3, 3)``. + rcut : float + Cutoff radius. + layout : GraphLayout or None + Edge-axis length policy. + with_csr : bool + Whether to construct destination/source CSR views. + canonicalize : bool + Whether to reorder every edge field into destination-major form. Implies + ``with_csr=True``. pair_excl Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` for model-level ``pair_exclude_types``. When given, @@ -151,7 +159,16 @@ def build_neighbor_graph_vesin( empty_S = torch.zeros((0, 3), dtype=torch.int64, device=dev) empty_nf = torch.zeros((0,), dtype=torch.int64, device=dev) return neighbor_graph_from_ijs( - empty_i, empty_i, empty_S, coord, box, empty_nf, nloc, layout=layout + empty_i, + empty_i, + empty_S, + coord, + box, + empty_nf, + nloc, + layout=layout, + with_csr=with_csr, + canonicalize=canonicalize, ) i_parts, j_parts, S_parts, nf_parts = [], [], [], [] @@ -199,9 +216,18 @@ def build_neighbor_graph_vesin( # out-of-cell (unwrapped) positions natively, so no normalize_coord is # needed and S is consistent with the original coords as searched. graph = neighbor_graph_from_ijs( - i_all, j_all, S_all, coord, box, nf_all, nloc, layout=layout + i_all, + j_all, + S_all, + coord, + box, + nf_all, + nloc, + layout=layout, ) if pair_excl is not None: at_flat = torch.as_tensor(atype, device=dev).reshape(-1) graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact) + if with_csr or canonicalize: + graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize) return graph diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 0c7e5a22ec..3daee763fb 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -229,15 +229,20 @@ def _build_single( ii, jj, ss = vesin_search_ijs( positions.detach(), cell if periodic else None, periodic, rcut, device ) - # ss is int64 from the helper; cast to float here for later ``ss @ box`` math. + # ss is int64 from the helper; cast to the coordinate dtype for the image sum. ss = ss.to(positions.dtype) # ghost atoms: neighbors reached through a non-zero periodic shift. Rebuild # their coordinates from the grad-carrying `positions`/`box` so autograd for - # forces/virials flows through the extended coordinates unchanged. + # forces/virials flows through the extended coordinates unchanged. The image + # offset ``S @ box`` is a broadcast multiply-reduce, not a matmul: the + # (n_ghost, 3) @ (3, 3) product otherwise dispatches to a full fp64 GEMM + # kernel whose length-3 contraction is catastrophically inefficient + # (measured ~2 ms vs ~30 us for the reduce on a 4k-atom cell), while + # remaining bit-identical. out_mask = torch.any(ss != 0, dim=1) out_idx = jj[out_mask] - out_coords = positions[out_idx] + ss[out_mask] @ box + out_coords = positions[out_idx] + (ss[out_mask][:, :, None] * box).sum(1) nghost = int(out_idx.shape[0]) extended_coord = torch.cat([positions, out_coords], dim=0) diff --git a/doc/install/install-lammps.md b/doc/install/install-lammps.md index 20ecdde56e..35278d6493 100644 --- a/doc/install/install-lammps.md +++ b/doc/install/install-lammps.md @@ -80,6 +80,16 @@ make -j4 make install ``` +To build the GPU-resident `pair_style deepmd/kk`, enable the LAMMPS Kokkos +package. The pair style uses the DeePMD C API: + +```bash +cmake -D PKG_KOKKOS=ON \ + -D LAMMPS_INSTALL_RPATH=ON -D BUILD_SHARED_LIBS=yes \ + -D CMAKE_INSTALL_PREFIX=${deepmd_root} \ + -D CMAKE_PREFIX_PATH=${deepmd_root} ../cmake +``` + If everything works fine, you will end up with an executable `${deepmd_root}/bin/lmp`. ```bash @@ -113,6 +123,12 @@ make -j4 make install ``` +Configure LAMMPS with `PKG_KOKKOS=ON`, then configure the DeePMD plugin with +`DEEPMD_LAMMPS_KOKKOS=ON` and `Kokkos_DIR` pointing to the Kokkos package +exported by the LAMMPS build. This registers the GPU-resident +`pair_style deepmd/kk`; the plugin and built-in integrations use the same C API +device-graph interface. + If everything works fine, you will end up with an executable `${deepmd_root}/bin/lmp`. ```bash diff --git a/pyproject.toml b/pyproject.toml index 56aa1edd1e..6216438ac4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -306,7 +306,7 @@ environment-pass = [ ] before-all = [ """if [ ! -z "${DP_PKG_NAME}" ]; then sed -i "s/name = \\"deepmd-kit\\"/name = \\"${DP_PKG_NAME}\\"/g" pyproject.toml; fi""", - """{ if [ -n "${CUDA_VERSION}" ] && [ "$(uname -m)" = "x86_64" ] ; then yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-nvcc-${CUDA_VERSION/./-} cuda-cudart-devel-${CUDA_VERSION/./-}; fi }""", + """{ if [ -n "${CUDA_VERSION}" ] && [ "$(uname -m)" = "x86_64" ] ; then yum config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-nvcc-${CUDA_VERSION/./-} cuda-cudart-devel-${CUDA_VERSION/./-}; fi }""", ] before-build = [ ] diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index c3f31e50dc..1a1102e9bd 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -356,6 +356,16 @@ if(ENABLE_PYTORCH AND NOT DEEPMD_C_ROOT) endif() endif() find_package(Torch REQUIRED) + # TorchConfig populates TORCH_CUDA_LIBRARIES only when PyTorch itself was + # built with CUDA. The fused graph-lower operators in op/pt include ATen CUDA + # headers and link libtorch_cuda, so they can be compiled only against a + # CUDA-enabled PyTorch: a CPU-only wheel ships the ATen/c10 CUDA headers but + # neither the generated c10/cuda/impl/cuda_cmake_macros.h nor libtorch_cuda. + if(TORCH_CUDA_LIBRARIES) + set(DEEPMD_TORCH_HAS_CUDA TRUE) + else() + set(DEEPMD_TORCH_HAS_CUDA FALSE) + endif() if(NOT Torch_VERSION VERSION_LESS "2.1.0") set_if_higher(CMAKE_CXX_STANDARD 17) elseif(NOT Torch_VERSION VERSION_LESS "1.5.0") diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 64585f58b6..eab09624df 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -1,5 +1,6 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #pragma once +#include #ifdef __cplusplus extern "C" { #else @@ -12,7 +13,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 27 +#define DP_C_API_VERSION 28 /** * @brief Neighbor list. @@ -358,6 +359,124 @@ extern void DP_DeepPotComputeNListf(DP_DeepPot* dp, float* atomic_energy, float* atomic_virial); +/** + * @brief Evaluate a device-resident edge graph with FP64 edge vectors. + * + * All ``d_*`` pointers reference accelerator memory on the model device. + * Frame and atomic parameters are host-resident arrays; pass NULL with size 0 + * to use the model defaults. + * + * @param[in] dp The DP 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_atom_virial Per-node virial, shape ``(nall_nodes, 9)``. + * @param[in] d_coord Per-node coordinates, shape ``(nall_nodes, 3)``. + * @param[in] d_atype Per-node atom types, shape ``(nall_nodes)``. + * @param[in] d_edge_index Destination-major ``[source, destination]`` edges, + * shape ``(2, nedge)``. + * @param[in] d_edge_vec Edge vectors, shape ``(nedge, 3)``. + * @param[in] nloc Number of owned local nodes. + * @param[in] nedge Number of physical edges. + * @param[in] fparam Host frame parameters. + * @param[in] fparam_size Number of values in ``fparam``. + * @param[in] aparam Host per-atom parameters. + * @param[in] aparam_size Number of values in ``aparam``. + * @param[in] nall_nodes Total local-plus-halo node count. + * @param[in] comm_nlist Communication metadata, or NULL. + * @since API version 28 + */ +extern void DP_DeepPotComputeEdgesGPU(DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + int nloc, + int nedge, + const double* fparam, + int64_t fparam_size, + const double* aparam, + int64_t aparam_size, + int nall_nodes, + const DP_Nlist* comm_nlist); + +/** + * @brief Evaluate a device-resident edge graph with FP32 edge vectors. + * + * The pointer and shape contract matches DP_DeepPotComputeEdgesGPU. + * + * @since API version 28 + */ +extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const float* d_edge_vec, + int nloc, + int nedge, + const double* fparam, + int64_t fparam_size, + const double* aparam, + int64_t aparam_size, + int nall_nodes, + const DP_Nlist* comm_nlist); + +/** + * @brief Evaluate a compact canonical graph on the model device. + * + * @param[in] dp The DP 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_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] nloc Number of owned local nodes. + * @param[in] nall_nodes Total local-plus-halo node count. + * @param[in] edge_storage Number of edge storage slots. + * @since API version 28 + */ +extern void DP_DeepPotComputeCanonicalGraphGPU( + DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const int64_t* d_atype, + const int64_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, + int nloc, + int nall_nodes, + int64_t edge_storage); + +/** + * @brief Query whether the loaded artifact supports device-edge inference. + * @since API version 28 + */ +extern bool DP_DeepPotSupportsDeviceEdgeInference(DP_DeepPot* dp); + +/** + * @brief Query whether device edge vectors use FP32. + * @since API version 28 + */ +extern bool DP_DeepPotUsesFP32EdgeVectors(DP_DeepPot* dp); + +/** + * @brief Query whether the compact canonical graph ABI is active. + * @since API version 28 + */ +extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); + /** * @brief Evaluate the energy, force and virial by using a DP. (double version) * @version 2 diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 69362165b8..939e32a381 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -958,7 +958,33 @@ struct InputNlist { recvproc, world, nprocs)) {}; - ~InputNlist() { DP_DeleteNlist(nl); }; + InputNlist(const InputNlist&) = delete; + InputNlist& operator=(const InputNlist&) = delete; + InputNlist(InputNlist&& other) noexcept + : nl(other.nl), + inum(other.inum), + ilist(other.ilist), + numneigh(other.numneigh), + firstneigh(other.firstneigh) { + other.nl = nullptr; + } + InputNlist& operator=(InputNlist&& other) noexcept { + if (this != &other) { + DP_DeleteNlist(nl); + nl = other.nl; + inum = other.inum; + ilist = other.ilist; + numneigh = other.numneigh; + firstneigh = other.firstneigh; + other.nl = nullptr; + } + return *this; + } + ~InputNlist() { + if (nl != nullptr) { + DP_DeleteNlist(nl); + } + }; /// @brief C API neighbor list. DP_Nlist* nl; /// @brief Number of core region atoms @@ -1193,6 +1219,112 @@ class DeepPot : public DeepBaseModel { return dchgspin; } + /** + * @brief Evaluate a device-resident edge graph with FP64 edge vectors. + * + * Edge, coordinate, type, and output pointers reside on the model device. + * Frame and atomic parameters remain host-resident because LAMMPS assembles + * them from computes, fixes, and pair-style settings. + */ + void compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes = 0, + const InputNlist* comm_nlist = nullptr) { + DP_DeepPotComputeEdgesGPU(dp, d_atom_energy, d_force, d_atom_virial, + d_coord, d_atype, d_edge_index, d_edge_vec, nloc, + nedge, fparam.empty() ? nullptr : fparam.data(), + static_cast(fparam.size()), + aparam.empty() ? nullptr : aparam.data(), + static_cast(aparam.size()), nall_nodes, + comm_nlist != nullptr ? comm_nlist->nl : nullptr); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + } + + /** + * @brief Evaluate a device-resident edge graph with FP32 edge vectors. + */ + void compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const float* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes = 0, + const InputNlist* comm_nlist = nullptr) { + DP_DeepPotComputeEdgesGPUFloat32( + dp, d_atom_energy, d_force, d_atom_virial, d_coord, d_atype, + d_edge_index, d_edge_vec, nloc, nedge, + fparam.empty() ? nullptr : fparam.data(), + static_cast(fparam.size()), + aparam.empty() ? nullptr : aparam.data(), + static_cast(aparam.size()), nall_nodes, + comm_nlist != nullptr ? comm_nlist->nl : nullptr); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + } + + /** + * @brief Evaluate a compact canonical graph on the model device. + */ + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const int64_t* d_atype, + const int64_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 int nloc, + const int nall_nodes, + const int64_t edge_storage) { + DP_DeepPotComputeCanonicalGraphGPU( + dp, d_atom_energy, d_force, d_atom_virial, d_atype, d_source, + d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, + nloc, nall_nodes, edge_storage); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + } + + /** + * @brief Query whether the loaded artifact supports device-edge inference. + */ + bool supports_device_edge_inference() const { + const bool result = DP_DeepPotSupportsDeviceEdgeInference(dp); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + return result; + } + + /** + * @brief Query whether the loaded artifact expects FP32 edge vectors. + */ + bool uses_fp32_edge_vectors() const { + const bool result = DP_DeepPotUsesFP32EdgeVectors(dp); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + return result; + } + + /** + * @brief Query whether the compact canonical graph ABI is active. + */ + bool uses_canonical_graph_inference() const { + const bool result = DP_DeepPotUsesCanonicalGraphInference(dp); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + return result; + } + /** * @brief Evaluate the energy, force and virial by using this DP. * @param[out] ener The system energy. diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index 3047688b8a..1c5623c40b 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -1557,6 +1557,51 @@ const char* string_to_char_exact(const std::string& str) { return buffer; } +static std::vector copy_optional_host_values(const double* values, + const int64_t size, + const char* name) { + if (size < 0 || (size > 0 && values == nullptr)) { + throw deepmd::deepmd_exception( + std::string(name) + + " requires a non-null pointer and non-negative size"); + } + return size > 0 ? std::vector(values, values + size) + : std::vector(); +} + +template +static void DP_DeepPotComputeEdgesGPU_variant(DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const EDGE_TYPE* d_edge_vec, + const int nloc, + const int nedge, + const double* fparam, + const int64_t fparam_size, + const double* aparam, + const int64_t aparam_size, + const int nall_nodes, + const DP_Nlist* comm_nlist) { + try { + const auto fparam_values = + copy_optional_host_values(fparam, fparam_size, "fparam"); + const auto aparam_values = + copy_optional_host_values(aparam, aparam_size, "aparam"); + const deepmd::InputNlist* communication = + comm_nlist != nullptr ? &comm_nlist->nl : nullptr; + dp->dp.compute_edges_gpu(d_atom_energy, d_force, d_atom_virial, d_coord, + d_atype, d_edge_index, d_edge_vec, nloc, nedge, + fparam_values, aparam_values, nall_nodes, + communication); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + } +} + extern "C" { const char* DP_NlistCheckOK(DP_Nlist* nlist) { @@ -1629,6 +1674,97 @@ void DP_DeepPotComputeNListf(DP_DeepPot* dp, force, virial, atomic_energy, atomic_virial); } +void DP_DeepPotComputeEdgesGPU(DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + const int nloc, + const int nedge, + const double* fparam, + const int64_t fparam_size, + const double* aparam, + const int64_t aparam_size, + const int nall_nodes, + const DP_Nlist* comm_nlist) { + DP_DeepPotComputeEdgesGPU_variant(dp, d_atom_energy, d_force, d_atom_virial, + d_coord, d_atype, d_edge_index, d_edge_vec, + nloc, nedge, fparam, fparam_size, aparam, + aparam_size, nall_nodes, comm_nlist); +} + +void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const float* d_edge_vec, + const int nloc, + const int nedge, + const double* fparam, + const int64_t fparam_size, + const double* aparam, + const int64_t aparam_size, + const int nall_nodes, + const DP_Nlist* comm_nlist) { + DP_DeepPotComputeEdgesGPU_variant(dp, d_atom_energy, d_force, d_atom_virial, + d_coord, d_atype, d_edge_index, d_edge_vec, + nloc, nedge, fparam, fparam_size, aparam, + aparam_size, nall_nodes, comm_nlist); +} + +void DP_DeepPotComputeCanonicalGraphGPU(DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const int64_t* d_atype, + const int64_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 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_atom_virial, d_atype, d_source, + d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, nloc, nall_nodes, edge_storage)); +} + +bool DP_DeepPotSupportsDeviceEdgeInference(DP_DeepPot* dp) { + try { + return dp->dp.supports_device_edge_inference(); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + return false; + } +} + +bool DP_DeepPotUsesFP32EdgeVectors(DP_DeepPot* dp) { + try { + return dp->dp.uses_fp32_edge_vectors(); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + return false; + } +} + +bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp) { + try { + return dp->dp.uses_canonical_graph_inference(); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + return false; + } +} + // multiple frames void DP_DeepPotCompute2(DP_DeepPot* dp, const int nframes, diff --git a/source/api_c/tests/test_deepmd_exception.cc b/source/api_c/tests/test_deepmd_exception.cc index 72dc1cd069..c281b563f0 100644 --- a/source/api_c/tests/test_deepmd_exception.cc +++ b/source/api_c/tests/test_deepmd_exception.cc @@ -8,6 +8,7 @@ #include #include #include +#include #include #include "deepmd.hpp" @@ -26,6 +27,17 @@ TEST(TestDeepmdException, deepmdexception_nofile) { deepmd::hpp::deepmd_exception); } +TEST(TestInputNlist, move_assignment_transfers_c_handle) { + deepmd::hpp::InputNlist source; + DP_Nlist* source_handle = source.nl; + deepmd::hpp::InputNlist target; + + target = std::move(source); + + EXPECT_EQ(target.nl, source_handle); + EXPECT_EQ(source.nl, nullptr); +} + TEST(TestChargeSpinValidation, exact_frame_values_use_input_storage) { std::vector charge_spin = {1.0, 2.0, 3.0, 4.0}; std::vector tiled; diff --git a/source/api_c/tests/test_deeppot_universal.cc b/source/api_c/tests/test_deeppot_universal.cc index d79c831295..1547ef4350 100644 --- a/source/api_c/tests/test_deeppot_universal.cc +++ b/source/api_c/tests/test_deeppot_universal.cc @@ -125,6 +125,9 @@ TEST_P(UniversalDeepPotCTest, Metadata) { EXPECT_EQ(DP_DeepPotGetDimAParam(dp), ref.dim_aparam); EXPECT_EQ(DP_DeepPotIsAParamNAll(dp), ref.aparam_nall); EXPECT_EQ(DP_DeepPotHasDefaultFParam(dp), ref.has_default_fparam); + EXPECT_FALSE(DP_DeepPotSupportsDeviceEdgeInference(dp)); + EXPECT_FALSE(DP_DeepPotUsesFP32EdgeVectors(dp)); + EXPECT_FALSE(DP_DeepPotUsesCanonicalGraphInference(dp)); const char* type_map = DP_DeepPotGetTypeMap(dp); EXPECT_STREQ(type_map, ref.type_map.c_str()); diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 4a197cb064..262c653713 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -308,13 +308,14 @@ class DeepPotBackend : public DeepBaseModelBackend { coord, atype, box, fparam, aparam, atomic); } /** - * @brief GPU-resident edge-input inference for the SeZM/DPA4 graph-form .pt2: - *given device edge tensors, write per-atom energy / force / virial back to - * the device output pointers. The PyTorch Exportable backend overrides this; - * every other backend inherits the throwing default. The signature is - * intentionally torch-free so the dispatcher stays backend-agnostic (no - * dynamic_cast into a PyTorch-heavy type, so ``libdeepmd_cc`` need not link - * PyTorch). + * @brief GPU-resident edge inference backend hook. + * + * Given device-resident edge tensors, write the per-atom energy, 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 DeepPot::compute_edges_gpu for + * the device pointer, graph, and communication contracts. */ virtual void compute_edges_gpu(double* d_atom_energy, double* d_force, @@ -325,6 +326,48 @@ class DeepPotBackend : public DeepBaseModelBackend { const double* d_edge_vec, const int nloc, const int nedge); + virtual void compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist); + virtual void compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const float* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist); + virtual void compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::int64_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 int nloc, + const int nall_nodes, + const std::int64_t edge_storage); + virtual bool supports_device_edge_inference() const; + virtual bool uses_fp32_edge_vectors() const; + virtual bool uses_canonical_graph_inference() const; }; /** @@ -683,7 +726,8 @@ class DeepPot : public DeepBaseModel { /** @} */ /** - * @brief Fully device-resident edge inference for single-domain SeZM/DPA4. + * @brief Fully device-resident inference for exported edge-input or + * graph-input .pt2 models. * * Forwards to the PyTorch Exportable (.pt2) backend's GPU edge path; raising * if the active backend is not ``DeepPotPTExpt``. All pointers reference GPU @@ -697,7 +741,8 @@ class DeepPot : public DeepBaseModel { * @param[out] d_atom_virial Per-atom virial, GPU [nloc * 9] row-major. * @param[in] d_coord Local coordinates, GPU [nloc * 3] row-major. * @param[in] d_atype Local atom types, GPU [nloc]. - * @param[in] d_edge_index Local edge graph, GPU [2 * nedge]. + * @param[in] d_edge_index Destination-major local edge graph, GPU + * [2 * nedge]. * @param[in] d_edge_vec Minimum-image bond vectors, GPU [nedge * 3]. * @param[in] nloc Number of local atoms. * @param[in] nedge Number of physical edges. @@ -712,6 +757,87 @@ class DeepPot : public DeepBaseModel { const int nloc, const int nedge); + /** + * @brief GPU-resident edge inference with runtime frame / atomic parameters. + * + * As the parameter-free overload, but ``fparam`` (global, ``dfparam`` values) + * and ``aparam`` (per-atom, ``nloc * daparam`` values) override the model's + * stored defaults. Empty vectors fall back to the stored default fparam and + * to no aparam, so the two overloads coincide. + * + * @param[in] fparam Host-resident runtime frame parameters, or empty for the + * model default. + * @param[in] aparam Host-resident runtime per-atom parameters (row-major + * [nloc, daparam]), or empty for none. + * @param[in] nall_nodes Total graph node count; 0 (or nloc) folds ghosts onto + * local owners (single domain), while nall_nodes > nloc keeps the extended + * (local + ghost) node set for a domain-decomposed run. + * @param[in] comm_nlist Communication neighbor list (send/recv swaps) for the + * extended node set. Required for a message-passing model under domain + * decomposition, where ghost features are exchanged across ranks inside the + * forward pass; nullptr otherwise. + */ + void compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes = 0, + const InputNlist* comm_nlist = nullptr); + + /** + * @brief Device-edge inference with FP32 edge vectors. + * + * Call this overload only when ``uses_fp32_edge_vectors()`` is true. + */ + void compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const float* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes = 0, + const InputNlist* comm_nlist = nullptr); + + /** + * @brief Whether the loaded artifact supports device-edge inference. + */ + bool supports_device_edge_inference() const; + + /** + * @brief Whether the loaded artifact expects FP32 device edge vectors. + */ + bool uses_fp32_edge_vectors() const; + + /** + * @brief Whether the loaded artifact uses the compact canonical graph ABI. + */ + bool uses_canonical_graph_inference() const; + + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::int64_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 int nloc, + const int nall_nodes, + const std::int64_t edge_storage); + int dim_chg_spin() const; protected: diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index acda36dad3..b03ab5290e 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -290,25 +290,46 @@ class DeepPotPTExpt : public DeepPotBackend { const std::vector& charge_spin, const bool atomic) override; + /** @brief Compatibility overload using stored frame and atomic parameters. */ + void compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + const int nloc, + const int nedge) override; + /** - * @brief Fully device-resident edge inference for single-domain SeZM/DPA4. + * @brief Fully device-resident inference for exported edge-input or + * graph-input .pt2 models. * * Runs the exported model directly on a GPU-built compact edge schema, * keeping coordinates, the edge graph and the outputs on the device. All * pointers reference GPU memory on the model's device. ``edge_index`` is the - * flattened [2, nedge] local edge graph (row 0 = neighbor/source, row 1 = + * flattened [2, nedge] edge graph (row 0 = neighbor/source, row 1 = * center/destination); ``edge_vec`` is the matching minimum-image bond vector * ``r_neighbor - r_center``. Outputs are written device-to-device. * * @param[out] d_atom_energy Per-atom energy, GPU [nloc]. - * @param[out] d_force Per-atom force, GPU [nloc * 3] row-major. - * @param[out] d_atom_virial Per-atom virial, GPU [nloc * 9] row-major. - * @param[in] d_coord Local coordinates, GPU [nloc * 3] row-major. - * @param[in] d_atype Local atom types, GPU [nloc]. - * @param[in] d_edge_index Local edge graph, GPU [2 * nedge]. + * @param[out] d_force Per-node force, GPU [nall_nodes * 3] row-major. + * @param[out] d_atom_virial Per-node virial, GPU [nall_nodes * 9] row-major. + * @param[in] d_coord Coordinates, GPU [nall_nodes * 3] row-major. + * @param[in] d_atype Atom types, GPU [nall_nodes]. + * @param[in] d_edge_index Destination-major [source, destination] edge graph, + * GPU [2 * nedge]. * @param[in] d_edge_vec Minimum-image bond vectors, GPU [nedge * 3]. * @param[in] nloc Number of local atoms. * @param[in] nedge Number of physical edges (dummy edges added internally). + * @param[in] fparam Host-resident frame parameters, or empty for the stored + * default. + * @param[in] aparam Host-resident per-atom parameters, or empty for none. + * @param[in] nall_nodes Graph node count: 0 (or nloc) folds ghosts onto local + * owners; nall_nodes > nloc keeps the extended (local + ghost) node set. + * @param[in] comm_nlist Communication neighbor list for the extended node + * set; required by a message-passing model under domain decomposition, + * nullptr otherwise. */ void compute_edges_gpu(double* d_atom_energy, double* d_force, @@ -318,9 +339,80 @@ class DeepPotPTExpt : public DeepPotBackend { const int* d_edge_index, const double* d_edge_vec, const int nloc, - const int nedge) override; + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist) override; + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::int64_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 int nloc, + const int nall_nodes, + const std::int64_t edge_storage) override; + /** + * @brief FP32 edge-vector overload for compressed graph artifacts. + * + * The model metadata must declare ``graph_edge_dtype == "float32"``. The + * remaining graph and output contracts are identical to the FP64 overload. + */ + void compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const float* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist) override; + /** + * @brief Whether the loaded artifact accepts the device-edge API. + */ + bool supports_device_edge_inference() const override; + /** + * @brief Whether the loaded graph artifact consumes FP32 edge vectors. + */ + bool uses_fp32_edge_vectors() const override; + bool uses_canonical_graph_inference() const override; private: + template + void compute_edges_gpu_impl(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const EDGE_TYPE* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist); + void compute_canonical_graph_gpu_impl( + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::int64_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 int nloc, + const int nall_nodes, + const std::int64_t edge_storage); bool inited; int ntypes; int dfparam; @@ -339,6 +431,8 @@ class DeepPotPTExpt : public DeepPotBackend { int nnei; // expected nlist nnei dimension (= sum(sel)) bool lower_input_is_edge_ = false; bool lower_input_is_graph_ = false; + bool lower_input_is_canonical_ = false; + bool graph_edge_fp32_ = false; NeighborListData nlist_data; at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) std::vector mapping_; // cached mapping vector (LAMMPS path) @@ -346,20 +440,16 @@ class DeepPotPTExpt : public DeepPotBackend { at::Tensor edge_index_tensor; // cached local edge graph (LAMMPS path) at::Tensor edge_index_ext_tensor; // cached extended edge graph (LAMMPS path) std::unique_ptr loader; - // Optional second AOTInductor artifact for the multi-rank GNN code - // path (Phase 4). Loaded only if the .pt2 metadata reports - // ``has_comm_artifact == true`` AND the model has GNN message - // passing. ``with_comm_tempfile_`` owns the extracted nested .pt2 - // for the lifetime of ``with_comm_loader``. + // Optional AOTInductor artifact for multi-rank message-passing inference, + // loaded only when the .pt2 metadata declares ``has_comm_artifact``. It is + // required when ghost features must be exchanged across ranks inside the + // forward pass. ``with_comm_tempfile_`` owns the extracted nested .pt2 for + // the lifetime of ``with_comm_loader``. bool has_comm_artifact_ = false; - // Whether the regular .pt2 graph consumes the mapping tensor for - // ghost-feature gather (true for any message-passing descriptor: - // DPA2/DPA3/hybrids; false for se_e2_a/DPA1/etc.). Mirrors the - // descriptor's ``has_message_passing()`` API; read from the - // ``has_message_passing`` metadata field. Defaults to false for - // pre-PR .pt2 archives that lack the field so non-GNN archives - // continue to work; GNN archives must be regenerated to opt into - // the fail-fast guard against the silent-corruption bug. + // Whether the regular .pt2 graph consumes the mapping tensor to gather ghost + // features, i.e. the descriptor performs message passing. Read from the + // ``has_message_passing`` metadata field; archives without it default to + // non-message-passing. bool has_message_passing_ = false; // Device-resident (ntypes+1)^2 model-level pair-type keep table, uploaded // ONCE in ``init`` from the ``pair_exclude_types`` metadata field (see @@ -441,28 +531,48 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& charge_spin); /** - * @brief Run a NeighborGraph-schema ``.pt2`` (lower_input_kind="graph"). + * @brief Run a graph-input ``.pt2`` artifact (lower_input_kind="graph"). * * Positional AOTI input order matches the Python export ABI: - * ``(atype, n_node, edge_index, edge_vec, edge_mask, [fparam], [aparam], - * [charge_spin])``. Unlike the edge schema there is no ``coord`` and no - * ``edge_scatter_index`` input; node count is carried by ``n_node`` and the - * geometry is fully described by ``edge_vec``. + * ``(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])``. * - * @param[in] atype Per-node local types, shape ``(N,)`` int64. - * @param[in] n_node Per-frame node count, shape ``(nf,)`` int64. + * @param[in] atype Per-node local-plus-halo types, shape ``(N,)`` int64. + * @param[in] n_node Per-frame total node count, shape ``(nf,)`` int64. + * @param[in] n_local Per-frame owned node count, shape ``(nf,)`` int64. * @param[in] edge_index Folded edge graph ``(2, E)`` int64 [src, dst]. * @param[in] edge_vec Edge vectors ``(E, 3)`` (neighbour - center). * @param[in] edge_mask Physical-edge mask ``(E,)`` bool. + * @param[in] destination_order Destination-grouped edge permutation ``(E,)``. + * @param[in] destination_row_ptr Destination CSR offsets ``(N + 1,)``. + * @param[in] source_order Source-grouped edge permutation ``(E,)``. + * @param[in] source_row_ptr Source CSR offsets ``(N + 1,)``. */ - std::vector run_model_graph(const torch::Tensor& atype, - const torch::Tensor& n_node, - const torch::Tensor& edge_index, - const torch::Tensor& edge_vec, - const torch::Tensor& edge_mask, - const torch::Tensor& fparam, - const torch::Tensor& aparam, - const torch::Tensor& charge_spin); + std::vector run_model_graph( + const torch::Tensor& atype, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& edge_index, + 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& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin); + + std::vector run_model_canonical_graph( + const torch::Tensor& atype, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& source, + const torch::Tensor& edge_vec, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_row_ptr, + const torch::Tensor& source_order); /** * @brief Run the with-comm .pt2 artifact with comm tensors appended. @@ -485,7 +595,7 @@ class DeepPotPTExpt : public DeepPotBackend { const std::vector& comm_tensors); /** - * @brief Run the with-comm edge (SeZM) ``.pt2`` artifact with comm tensors. + * @brief Run the with-comm edge-input ``.pt2`` artifact with comm tensors. * * The edge schema indexes the extended node set, so ``edge_index`` and * ``edge_scatter_index`` coincide. ``atype`` carries owned atoms (fitting, diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index eec3fdfbd2..12bd9754b5 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -371,28 +371,144 @@ inline EdgeTensorPack compactEdgeTensors(const torch::Tensor& edge_index, struct GraphTensorPack { torch::Tensor atype; torch::Tensor n_node; + torch::Tensor n_local; torch::Tensor edge_index; torch::Tensor edge_vec; torch::Tensor edge_mask; + torch::Tensor destination_order; + torch::Tensor destination_row_ptr; + torch::Tensor source_order; + torch::Tensor source_row_ptr; }; +struct CanonicalGraphTensorPack { + torch::Tensor 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; +}; + +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); + auto source = torch::zeros({storage_count}, + graph.edge_index.options().dtype(torch::kInt64)); + 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)); + 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)); + 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)); + } + return {graph.atype, + graph.n_node, + graph.n_local, + source, + edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + source_order}; +} + +/** + * @brief Build destination/source CSR views of an edge pack. + * + * Both views store permutations into the original edge payload. Masked + * padding entries form the suffix of each permutation. + * + * ``destination_sorted`` requires a destination-grouped real-edge prefix and + * a masked suffix. The destination permutation is the identity in this layout. + */ +inline void buildGraphCSR(GraphTensorPack& pack, + const std::int64_t node_count, + const bool destination_sorted = false) { + const auto real_index = torch::nonzero(pack.edge_mask).reshape({-1}); + const auto padding_index = + torch::nonzero(torch::logical_not(pack.edge_mask)).reshape({-1}); + const auto real_destination = + pack.edge_index.select(0, 1).index_select(0, real_index); + const auto real_source = + pack.edge_index.select(0, 0).index_select(0, real_index); + const auto destination_counts = + torch::bincount(real_destination, {}, node_count); + const auto source_counts = torch::bincount(real_source, {}, node_count); + const auto zero = torch::zeros({1}, destination_counts.options()); + pack.destination_row_ptr = + torch::cat({zero, torch::cumsum(destination_counts, 0)}) + .to(torch::kInt64) + .contiguous(); + pack.source_row_ptr = torch::cat({zero, torch::cumsum(source_counts, 0)}) + .to(torch::kInt64) + .contiguous(); + const auto real_source_order = torch::argsort(real_source, 0, false); + if (destination_sorted) { + pack.destination_order = + torch::arange(pack.edge_index.size(1), real_index.options()); + } else { + const auto real_destination_order = + torch::argsort(real_destination, 0, false); + pack.destination_order = + torch::cat( + {real_index.index_select(0, real_destination_order), padding_index}) + .contiguous(); + } + pack.source_order = + torch::cat({real_index.index_select(0, real_source_order), padding_index}) + .contiguous(); +} + +inline void buildGraphCSR(GraphTensorPack& pack) { + buildGraphCSR(pack, pack.atype.size(0)); +} + +/** + * @brief Reorder an edge payload into destination-major form and build CSR. + * + * The same permutation is applied to every edge field. Masked entries move to + * the suffix, and the resulting destination permutation is the identity. + */ +inline void canonicalizeGraphPayload(GraphTensorPack& pack, + const std::int64_t node_count) { + const auto destination = pack.edge_index.select(0, 1); + const auto padding_node = torch::full_like(destination, node_count); + const auto destination_key = + torch::where(pack.edge_mask, destination, padding_node); + const auto order = torch::argsort(destination_key, /*stable=*/true, 0, false); + pack.edge_index = pack.edge_index.index_select(1, order).contiguous(); + pack.edge_vec = pack.edge_vec.index_select(0, order).contiguous(); + pack.edge_mask = pack.edge_mask.index_select(0, order).contiguous(); + buildGraphCSR(pack, node_count, /*destination_sorted=*/true); +} + /** * @brief Build NeighborGraph input tensors from a host neighbor list * (single-rank, dynamic edge axis). * * Mirrors the edge schema but drops ``coord``/``edge_scatter_index`` and adds - * ``n_node``. Edge construction is delegated to the existing - * ``createEdgeTensors``/``compactEdgeTensors`` helpers (same rcut filter, - * variable edge count and two masked dummy edges that keep the dynamic edge - * dimension non-empty); the wrapper then (a) drops the extended scatter index, - * (b) emits ``n_node = [nloc]`` for the single frame, and (c) sets the node - * types from the local slice of ``atype_ext``. + * ``n_node``. Edge construction is delegated to + * ``createEdgeTensors``/``compactEdgeTensors``; the wrapper emits + * ``n_node = [nloc]`` for the single frame and sets node types from the local + * slice of ``atype_ext``. * * @param nlist Neighbor-list rows (local idx into the extended set). * @param coord Extended coordinates shaped as nall x 3. * @param atype_ext Extended atom types, length nall. Node types are taken from * the extended types (NOT ``atype[mapping]``); for single-rank ghost-free - * this is just ``atype_ext[0:nloc]``, while multi-rank (B3) passes the halo + * this is just ``atype_ext[0:nloc]``; an extended graph also carries the halo * types directly. * @param mapping Extended-to-local atom map, length nall. * @param nloc Number of local atoms. @@ -405,7 +521,7 @@ struct GraphTensorPack { * owners (single-rank, ``N == nloc``, ``n_node = [nloc]``, node types from * ``atype_ext[0:nloc]``) or kept as distinct extended nodes (multi-rank, * ``N == nall``, ``n_node = [nall]``, node types from the full ``atype_ext`` - * including the real halo types — the #5583 invariant). In the multi-rank + * including the real halo types). In the multi-rank * case ``edge_index`` indexes the extended atoms directly, so ghost reaction * forces land on the ghost rows and are folded to their owners by LAMMPS * reverse-comm (no with-comm artifact / no border_op — dpa1 is non-MP). @@ -429,7 +545,7 @@ inline GraphTensorPack buildGraphTensors( // when false, edge_index indexes the extended atoms directly (multi-rank). // edge_index_ext always keeps extended indices for the on-device geometry // recompute. - const EdgeTensorPack topo = + const EdgeTensorPack topology = createEdgeTensors(nlist, coord, mapping, nloc, nall, device, /*with_geometry=*/false, row_centers, fold_to_local); @@ -443,7 +559,7 @@ inline GraphTensorPack buildGraphTensors( .clone() .to(device); const EdgeTensorPack edges = compactEdgeTensors( - topo.edge_index, topo.edge_index_ext, coord_tensor, rcut); + topology.edge_index, topology.edge_index_ext, coord_tensor, rcut); GraphTensorPack pack; pack.edge_index = edges.edge_index; // (2, E): local-folded or extended @@ -453,6 +569,7 @@ inline GraphTensorPack buildGraphTensors( // (ghosts are distinct nodes whose features come from their real halo types). const std::int64_t n_node_count = fold_to_local ? nloc : nall; pack.n_node = torch::full({1}, n_node_count, int_options).to(device); + pack.n_local = torch::full({1}, nloc, int_options).to(device); // Node types from the extended types (NOT atype[mapping]): the local slice // for single-rank, the full extended set (incl. real halo types) for // multi-rank. @@ -588,7 +705,7 @@ inline torch::Tensor applyPairExclusionNlist( * rows already carry the folded ghost contributions, so zero ghosts avoid * double counting (and keep LAMMPS reverse-comm correct). * - * **Single-rank only.** Multi-rank inference (B3.2) must NOT call this + * **Single-rank only.** Multi-rank inference must not call this * function: ghost/halo forces are real cross-rank contributions that must be * returned as-is and folded back via reverse-comm rather than being zeroed. * Calling this function on a multi-rank result would silently zero those forces @@ -610,8 +727,7 @@ inline void remap_graph_outputs_to_dense_keys( const bool single_rank = true) { if (!single_rank) { throw deepmd::deepmd_exception( - "remap_graph_outputs_to_dense_keys is single-rank-only; multi-rank " - "uses the extended-region reverse-comm fold (PR-B3.2)"); + "remap_graph_outputs_to_dense_keys requires single-rank graph input"); } using torch::indexing::Slice; const std::int64_t nf = 1; diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index d2b9d773ea..4888e0b387 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -582,12 +582,116 @@ void DeepPotBackend::compute_edges_gpu(double* d_atom_energy, (void)d_edge_vec; (void)nloc; (void)nedge; + throw deepmd::deepmd_exception( + "compute_edges_gpu (GPU-resident edge inference) is only supported by " + "the PyTorch Exportable (.pt2) backend."); +} + +void DeepPotBackend::compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist) { + if (fparam.empty() && aparam.empty() && + (nall_nodes == 0 || nall_nodes == nloc) && comm_nlist == nullptr) { + compute_edges_gpu(d_atom_energy, d_force, d_atom_virial, d_coord, d_atype, + d_edge_index, d_edge_vec, nloc, nedge); + return; + } + (void)d_atom_energy; + (void)d_force; + (void)d_atom_virial; + (void)d_coord; + (void)d_atype; + (void)d_edge_index; + (void)d_edge_vec; + (void)comm_nlist; + (void)nloc; + (void)nedge; + (void)fparam; + (void)aparam; + (void)nall_nodes; throw deepmd::deepmd_exception( "compute_edges_gpu (GPU-resident edge inference) is only supported by " "the " "PyTorch Exportable (.pt2) backend."); } +void DeepPotBackend::compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const float* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist) { + (void)d_atom_energy; + (void)d_force; + (void)d_atom_virial; + (void)d_coord; + (void)d_atype; + (void)d_edge_index; + (void)d_edge_vec; + (void)comm_nlist; + (void)nloc; + (void)nedge; + (void)fparam; + (void)aparam; + (void)nall_nodes; + throw deepmd::deepmd_exception( + "compute_edges_gpu with float32 edge vectors is only supported by a " + "compatible PyTorch Exportable (.pt2) backend."); +} + +void DeepPotBackend::compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::int64_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 int nloc, + const int nall_nodes, + const std::int64_t edge_storage) { + (void)d_atom_energy; + (void)d_force; + (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)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 DeepPotBackend::uses_fp32_edge_vectors() const { return false; } + +bool DeepPotBackend::supports_device_edge_inference() const { return false; } + +bool DeepPotBackend::uses_canonical_graph_inference() const { return false; } + void DeepPot::compute_edges_gpu(double* d_atom_energy, double* d_force, double* d_atom_virial, @@ -597,16 +701,81 @@ void DeepPot::compute_edges_gpu(double* d_atom_energy, const double* d_edge_vec, const int nloc, const int nedge) { - // Polymorphic dispatch to the loaded backend: the PyTorch Exportable backend - // overrides ``compute_edges_gpu``; other backends inherit the throwing - // default. This replaces a ``dynamic_cast`` into the PyTorch-heavy - // ``DeepPotPTExpt``, which the "load backends as plugins" refactor made - // uncompilable in the backend-agnostic ``libdeepmd_cc`` (it does not link - // PyTorch), so the cast branch was always stubbed out. dp->compute_edges_gpu(d_atom_energy, d_force, d_atom_virial, d_coord, d_atype, d_edge_index, d_edge_vec, nloc, nedge); } +void DeepPot::compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist) { + // Backend-agnostic dispatch: backends that implement device edge inference + // override ``compute_edges_gpu``, 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_edges_gpu(d_atom_energy, d_force, d_atom_virial, d_coord, d_atype, + d_edge_index, d_edge_vec, nloc, nedge, fparam, aparam, + nall_nodes, comm_nlist); +} + +void DeepPot::compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const float* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist) { + dp->compute_edges_gpu(d_atom_energy, d_force, d_atom_virial, d_coord, d_atype, + d_edge_index, d_edge_vec, nloc, nedge, fparam, aparam, + nall_nodes, comm_nlist); +} + +void DeepPot::compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::int64_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 int nloc, + const int nall_nodes, + const std::int64_t edge_storage) { + dp->compute_canonical_graph_gpu( + d_atom_energy, d_force, d_atom_virial, d_atype, d_source, d_edge_vec, + d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, + edge_storage); +} + +bool DeepPot::uses_fp32_edge_vectors() const { + return dp->uses_fp32_edge_vectors(); +} + +bool DeepPot::supports_device_edge_inference() const { + return dp->supports_device_edge_inference(); +} + +bool DeepPot::uses_canonical_graph_inference() const { + return dp->uses_canonical_graph_inference(); +} + int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); } DeepPotModelDevi::DeepPotModelDevi() { diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index a60bd681e4..8a81edd284 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -2,6 +2,7 @@ #include "DeepPotPTExpt.h" #if defined(BUILD_PYTORCH) && BUILD_PT_EXPT +#include #include #include @@ -11,6 +12,7 @@ #include #include #include +#include #include "SimulationRegion.h" #include "common.h" @@ -26,6 +28,82 @@ 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 +} + +template +at::Tensor make_aparam_tensor(const std::vector& aparam, + std::int64_t nloc, + std::int64_t daparam, + const torch::Device& device) { + const auto float64_options = torch::TensorOptions().dtype(torch::kFloat64); + if (daparam == 0) { + if (!aparam.empty()) { + throw deepmd::deepmd_exception( + "The model does not accept atomic parameters, but aparam is not " + "empty."); + } + return torch::zeros({0}, float64_options).to(device); + } + + const std::int64_t expected = nloc * daparam; + if (static_cast(aparam.size()) != expected) { + throw deepmd::deepmd_exception( + "aparam has " + std::to_string(aparam.size()) + + " values, but the model expects nloc * daparam = " + + std::to_string(expected) + "."); + } + if (expected == 0) { + return torch::zeros({1, nloc, daparam}, float64_options).to(device); + } + + const auto input_options = + std::is_same::value + ? torch::TensorOptions().dtype(torch::kFloat32) + : torch::TensorOptions().dtype(torch::kFloat64); + return torch::from_blob(const_cast(aparam.data()), + {1, nloc, daparam}, input_options) + .to(torch::kFloat64) + .to(device); +} + +struct FrameParameterLayout { + std::size_t frame_size; + bool broadcast; +}; + +FrameParameterLayout resolve_frame_parameter_layout(const char* name, + std::size_t size, + int nframes, + std::size_t frame_size, + bool allow_empty) { + if (size == 0 && allow_empty) { + return {0, true}; + } + const std::size_t full_size = static_cast(nframes) * frame_size; + if (size == frame_size) { + return {frame_size, true}; + } + if (size == full_size) { + return {frame_size, false}; + } + throw deepmd::deepmd_exception( + std::string(name) + " has " + std::to_string(size) + + " values, but the model expects " + std::to_string(frame_size) + + " for frame broadcasting or " + std::to_string(full_size) + + " for frame-specific values."); +} + +} // namespace + void DeepPotPTExpt::translate_error(std::function f) { try { f(); @@ -67,8 +145,8 @@ void DeepPotPTExpt::init(const std::string& model, // Load libdeepmd_op_pt.so so its TORCH_LIBRARY_FRAGMENT entries // (deepmd::*, deepmd_export::*) are visible to torch's dispatcher - // before the AOTI module loads. Without this, multi-rank GNN .pt2 - // archives fail at pair_style time with + // before the AOTI module loads. Without this, multi-rank message-passing + // .pt2 archives fail at pair_style time with // ``Could not find schema for deepmd_export::border_op``. deepmd::load_op_library(deepmd::DPBackend::PyTorchExportable); @@ -157,9 +235,27 @@ 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"; } else { lower_input_is_edge_ = false; lower_input_is_graph_ = false; + lower_input_is_canonical_ = false; + } + graph_edge_fp32_ = false; + if (metadata.obj_val.count("graph_edge_dtype")) { + const std::string graph_edge_dtype = + metadata["graph_edge_dtype"].as_string(); + if (graph_edge_dtype != "float32" && graph_edge_dtype != "float64") { + throw deepmd::deepmd_exception( + "metadata graph_edge_dtype must be 'float32' or 'float64'."); + } + graph_edge_fp32_ = graph_edge_dtype == "float32"; + } + if (lower_input_is_canonical_) { + if (!graph_edge_fp32_) { + throw deepmd::deepmd_exception( + "compact canonical graph artifacts require float32 edge vectors."); + } } type_map.clear(); @@ -179,26 +275,18 @@ void DeepPotPTExpt::init(const std::string& model, gpu_enabled ? static_cast(gpu_id) : static_cast(-1)); - // Phase 4: load the optional with-comm artifact for multi-rank GNN - // inference. Pre-Phase-3 .pt2 files lack ``has_comm_artifact``; - // default to false so old artifacts keep working. If the metadata - // flag is set but the nested artifact fails to extract or compile, - // keep ``has_comm_artifact_=true`` and let single-rank dispatch - // continue working; multi-rank dispatch then fails fast at - // ``run_model_with_comm()`` rather than silently dropping the MPI - // exchange and producing wrong results. + // The optional with-comm artifact drives multi-rank message-passing + // inference. Archives without ``has_comm_artifact`` default to false. When + // the flag is set but the nested artifact fails to extract or compile, + // ``has_comm_artifact_`` stays true so single-rank dispatch keeps working + // while multi-rank dispatch raises at ``run_model_with_comm`` rather than + // silently dropping the ghost exchange. has_comm_artifact_ = metadata.obj_val.count("has_comm_artifact") && metadata["has_comm_artifact"].as_bool(); - // Whether the regular .pt2 graph consumes ``mapping`` for ghost-atom - // feature gather. Mirrors the descriptor's ``has_message_passing()`` - // API: true for message-passing descriptors (DPA2, DPA3, hybrids - // over those), false for non-message-passing descriptors (se_e2_a, - // DPA1, etc.). Pre-PR .pt2 archives lack this field; default to - // false so they retain their previous behaviour (non-GNN archives - // continue to work; GNN archives that had the original - // silent-corruption bug must be regenerated to opt into the fail- - // fast guard). All in-tree fixtures are regenerated by the gen - // scripts and carry the explicit value. + // Whether the regular .pt2 graph consumes ``mapping`` to gather ghost + // features, i.e. the descriptor performs message passing. Read from the + // ``has_message_passing`` metadata field; archives without it default to + // non-message-passing. has_message_passing_ = metadata.obj_val.count("has_message_passing") && metadata["has_message_passing"].as_bool(); @@ -328,16 +416,22 @@ std::vector DeepPotPTExpt::run_model_edges( std::vector DeepPotPTExpt::run_model_graph( const torch::Tensor& atype, const torch::Tensor& n_node, + const torch::Tensor& n_local, const torch::Tensor& edge_index, 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& fparam, const torch::Tensor& aparam, const torch::Tensor& charge_spin) { - // NeighborGraph ABI: (atype, n_node, edge_index, edge_vec, edge_mask, - // [fparam], [aparam], [charge_spin]). No coord, no edge_scatter_index. - std::vector inputs = {atype, n_node, edge_index, edge_vec, - edge_mask}; + // Graph-input ABI: original edge payload plus destination/source CSR views. + std::vector inputs = { + atype, n_node, n_local, edge_index, + edge_vec, edge_mask, destination_order, destination_row_ptr, + source_order, source_row_ptr}; if (dfparam > 0) { inputs.push_back(fparam); } @@ -350,6 +444,19 @@ std::vector DeepPotPTExpt::run_model_graph( return loader->run(inputs); } +std::vector DeepPotPTExpt::run_model_canonical_graph( + const torch::Tensor& atype, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& source, + const torch::Tensor& edge_vec, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_row_ptr, + const torch::Tensor& source_order) { + return loader->run({atype, n_node, n_local, source, edge_vec, + destination_row_ptr, source_row_ptr, source_order}); +} + std::vector DeepPotPTExpt::run_model_with_comm( const torch::Tensor& coord, const torch::Tensor& atype, @@ -361,12 +468,9 @@ std::vector DeepPotPTExpt::run_model_with_comm( const std::vector& comm_tensors) { if (!with_comm_loader) { throw deepmd::deepmd_exception( - "run_model_with_comm called but the with-comm artifact is not " - "available. Either the .pt2 file has no with-comm artifact compiled " - "(programming error: the caller should check has_comm_artifact_ " - "before invoking this path), or the artifact was present in the " - ".pt2 metadata but failed to load at init time (see earlier stderr " - "log). Multi-rank LAMMPS requires a working with-comm artifact."); + "The with-comm artifact is unavailable for this model. Multi-rank " + "message-passing inference requires a with-comm artifact that is both " + "present in the .pt2 metadata and successfully loaded."); } if (comm_tensors.size() != 8) { throw deepmd::deepmd_exception( @@ -405,12 +509,9 @@ std::vector DeepPotPTExpt::run_model_edges_with_comm( const std::vector& comm_tensors) { if (!with_comm_loader) { throw deepmd::deepmd_exception( - "run_model_edges_with_comm called but the with-comm artifact is not " - "available. Either the .pt2 file has no with-comm artifact compiled " - "(programming error: the caller should check has_comm_artifact_ " - "before invoking this path), or the artifact was present in the " - ".pt2 metadata but failed to load at init time (see earlier stderr " - "log). Multi-rank LAMMPS requires a working with-comm artifact."); + "The with-comm artifact is unavailable for this model. Multi-rank " + "message-passing inference requires a with-comm artifact that is both " + "present in the .pt2 metadata and successfully loaded."); } if (comm_tensors.size() != 8) { throw deepmd::deepmd_exception( @@ -521,10 +622,9 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // Dispatch decision: use the with-comm artifact when LAMMPS is running // multi-rank. ``lmp_list.nprocs > 1`` is the direct predicate; // LAMMPS pair styles populate it by passing ``comm->nprocs`` to the - // ``InputNlist`` constructor. Earlier drafts used ``nswap > 0`` as a - // proxy, but that breaks for ``atom_style spin`` (which emits - // nswap > 0 even in single-rank to propagate PBC ghost spins). - // ``nprocs`` is unambiguous. + // ``InputNlist`` constructor. ``nprocs`` is the unambiguous predicate; + // ``nswap`` is nonzero even single-rank for ``atom_style spin`` (periodic + // ghost-spin propagation), so it cannot serve as a multi-rank proxy. // // The regular artifact uses ``mapping`` to gather ghost-atom features // from local-atom embeddings (``index_select(node_ebd[1, nloc, dim], @@ -536,32 +636,15 @@ 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; - // 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 - // distinct nodes whose features come from their real halo types). No - // with-comm artifact / no border_op is needed; ghost reaction forces are - // folded to their owners by LAMMPS reverse-comm. Handled below. - // - message-passing graph (DPA2/DPA3, PR-G): would need a with-comm graph - // artifact for cross-rank ghost-feature exchange — not yet supported. - // Fail fast before building any tensors so callers get a clear message - // instead of a wrong answer. if (lower_input_is_graph_ && multi_rank && has_message_passing_) { throw deepmd::deepmd_exception( - "Multi-rank message-passing graph (NeighborGraph) .pt2 inference is " - "not yet supported (PR-G). Non-message-passing graph models (e.g. " - "dpa1) run multi-rank on the extended-region single-rank artifact; " - "for message-passing models run single-rank, or use a dense/edge " - ".pt2 for multi-rank LAMMPS."); - } - // Decision matrix (see PR #5450 description): - // non-GNN model (has_message_passing_ == false): regular path is - // always safe. - // nghost == 0 (NoPbc, isolated cluster): always safe. - // GNN model, multi-rank: requires has_comm_artifact_ (cell C-mr / D-mr) - // else fail-fast (cell B-mr) - // GNN model, single-rank: requires atom_map_present (cell A / C) - // else fail-fast (cell B / D) + "Multi-rank message-passing graph-input inference requires an " + "edge-input model with its with-comm artifact."); + } + // Dispatch guards for message-passing models: + // 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 (multi_rank && !has_comm_artifact_) { throw deepmd::deepmd_exception( @@ -592,10 +675,9 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // is consumed every step by the dense ``run_model`` (ghost-feature gather); // the ``mapping_`` vector is read only here at ago==0 -- to build that // tensor and, for the edge/graph paths, to fold ghost neighbours onto their - // local owners inside ``createEdgeTensors``. (The graph path used to read - // ``mapping_`` every step via a per-step ``buildGraphTensors``; it now - // caches the topology at ago==0 like the edge/dense paths, so no per-step - // read.) + // local owners inside ``createEdgeTensors``. The edge and graph paths + // cache this topology at ago==0, so no ``mapping_`` read occurs on the + // per-step path. if (lmp_list.mapping) { mapping_.resize(nall_real); for (int ii = 0; ii < nall_real; ii++) { @@ -641,14 +723,10 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, /*fold_to_local=*/!use_with_comm); edge_index_tensor = edge_tensors.edge_index; edge_index_ext_tensor = edge_tensors.edge_index_ext; - } else if (lower_input_is_graph_) { - // Cache only the real skin topology, exactly like the edge path: the - // geometry (edge_vec) + rcut filter are recomputed on-device every step - // by compactEdgeTensors, so the O(E) host loop + H2D copy in - // createEdgeTensors runs ONLY on a LAMMPS nlist rebuild (ago==0), not - // every step. Single-rank folds ghosts onto local owners - // (fold_to_local=true); non-MP multi-rank keeps the extended region - // (fold_to_local=false) so ghost forces reverse-comm to their owners. + } else if (lower_input_is_graph_ || lower_input_is_canonical_) { + // Cache the skin topology. Single-rank folds ghosts onto local owners; + // non-message-passing multi-rank keeps the extended region so ghost + // forces reverse-comm to their owners. const auto edge_tensors = createEdgeTensors( nlist_data.jlist, dcoord, mapping_, nloc, nall_real, device, /*with_geometry=*/false, /*row_centers=*/&nlist_data.ilist, @@ -690,18 +768,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, fparam_tensor = torch::zeros({0}, options).to(device); } - at::Tensor aparam_tensor; - if (!aparam_.empty()) { - aparam_tensor = - torch::from_blob( - const_cast(aparam_.data()), - {1, nloc, static_cast(aparam_.size()) / nloc}, - valuetype_options) - .to(torch::kFloat64) - .to(device); - } else { - aparam_tensor = torch::zeros({0}, options).to(device); - } + 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. @@ -752,14 +819,10 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, std::vector> remapped_sendlist; std::vector remapped_sendlist_ptrs; std::vector remapped_sendnum, remapped_recvnum; - // Empty-subdomain phantom padding (edge with-comm path only): a rank that - // owns zero local atoms would feed nloc==0 / nedge==0 into the with-comm - // artifact, which is traced with nloc_min=1 / nedge_min=2 and may be lowered - // by inductor under an even stricter nloc>=2 assumption -- so a 0-atom rank - // can SIGFPE or silently corrupt instead of throwing. Mirror the spin - // phantom-atom workaround (DeepSpinPTExpt): run the graph with two phantom - // local atoms and two masked self-edges, then report the empty rank's true - // (zero) contribution. Setting the comm ``nlocal`` to ``phantom_n`` makes + // Empty-subdomain padding keeps every rank inside the exported lower bounds + // and participating in communication. Two phantom local atoms and two masked + // self-edges contribute zero physical output. Setting comm ``nlocal`` to + // ``phantom_n`` makes // border_op write received ghost features past the phantom slots, so // collective communication stays in lockstep with non-empty ranks. // Gated on ``use_with_comm`` so the strip-back below never touches outputs @@ -841,7 +904,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, ph_edge_index, ph_edge_mask, fparam_tensor, ph_aparam, charge_spin_tensor, comm_tensors); } else { - // SeZM edge schema: edge_index already indexes the extended node set + // Edge-input schema: edge_index already indexes the extended node set // (fold_to_local=false above), so it doubles as the force-scatter // index. Ghost node features are exchanged between blocks inside the // with-comm graph via border_op; the local atom types feed fitting @@ -881,12 +944,11 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, edge_tensors.edge_index, edge_tensors.edge_vec, edge_tensors.edge_index_ext, edge_tensors.edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor); - } else if (lower_input_is_graph_) { + } else if (lower_input_is_graph_ || lower_input_is_canonical_) { if (nall_real == 0) { // Truly-empty rank (no local atoms AND no ghosts): the graph would emit - // N == 0 nodes, and edge_force_virial's ``edge_index % node_capacity`` - // would divide by zero (SIGFPE) -- it also violates the exported - // ``Dim("n_node_total", min=1)``. Such a rank contributes nothing, so + // N == 0 nodes, which violates the exported + // ``Dim("n_node_total", min=1)``. Such a rank contributes nothing, so // fill zero outputs and return instead of running the model. (The // tested ``nloc == 0`` empty-subdomain case has ``nall_real > 0`` -- // ghosts within rcut -- so it still runs the model normally.) @@ -903,34 +965,52 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, } return; } - // NeighborGraph schema: recompute geometry + rcut filter on-device from - // the cached skin topology (edge_index[_ext]_tensor built at ago==0), - // then assemble the cheap node tensors. Mirrors the edge path -- no - // per-step host rebuild / H2D copy. Single-rank folds ghosts onto local - // owners (N == nloc); multi-rank (non-MP only — the fail-fast above - // blocks MP graph multi-rank) keeps the extended region (N == nall_real, - // node types from the real halo types) so LAMMPS reverse-comm folds ghost - // forces back. The node types come from the on-device extended - // atype_Tensor slice (== atype_ext[0:N]); n_node is a 1-element tensor. + // Compact the cached skin topology to the current model cutoff. + // Single-rank folds ghosts onto local owners (N == nloc); + // non-message-passing multi-rank keeps the extended region + // (N == nall_real) so reverse communication folds ghost forces back. const auto edge_tensors = compactEdgeTensors(edge_index_tensor, edge_index_ext_tensor, coord_Tensor, static_cast(rcut)); const std::int64_t n_node_count = multi_rank ? nall_real : nloc; at::Tensor n_node_tensor = torch::full({1}, n_node_count, int_option).to(device); + at::Tensor n_local_tensor = torch::full({1}, nloc, int_option).to(device); at::Tensor node_atype = atype_Tensor.slice(1, 0, n_node_count).reshape({n_node_count}); - // Model-level pair exclusion is a BUILD-time transform (decision - // #18/A4): the exported graph lower consumes a pre-excluded edge_mask - // and never re-applies it; this is the single application site on the - // C++ graph route. - const at::Tensor graph_edge_mask = deepmd::applyPairExclusion( + at::Tensor graph_aparam = aparam_tensor; + if (daparam > 0 && n_node_count > nloc) { + graph_aparam = torch::cat( + {aparam_tensor, torch::zeros({1, n_node_count - nloc, daparam}, + aparam_tensor.options())}, + 1); + } + GraphTensorPack graph_pack; + graph_pack.atype = node_atype; + graph_pack.n_node = n_node_tensor; + graph_pack.n_local = n_local_tensor; + graph_pack.edge_index = edge_tensors.edge_index; + graph_pack.edge_vec = graph_edge_fp32_ + ? edge_tensors.edge_vec.to(torch::kFloat32) + : edge_tensors.edge_vec; + graph_pack.edge_mask = deepmd::applyPairExclusion( edge_tensors.edge_index, edge_tensors.edge_mask, node_atype, pair_exclude_table_, ntypes); - flat_outputs = - run_model_graph(node_atype, n_node_tensor, edge_tensors.edge_index, - edge_tensors.edge_vec, graph_edge_mask, fparam_tensor, - aparam_tensor, charge_spin_tensor); + canonicalizeGraphPayload(graph_pack, n_node_count); + if (lower_input_is_canonical_) { + const auto compact = compactCanonicalGraph(graph_pack); + flat_outputs = run_model_canonical_graph( + compact.atype, compact.n_node, compact.n_local, compact.source, + compact.edge_vec, compact.destination_row_ptr, + compact.source_row_ptr, compact.source_order); + } else { + flat_outputs = run_model_graph( + node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index, + graph_pack.edge_vec, graph_pack.edge_mask, + graph_pack.destination_order, graph_pack.destination_row_ptr, + graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, + graph_aparam, charge_spin_tensor); + } } else { // Model-level pair exclusion is a BUILD-time transform (decision // #18/A4): the exported dense lower consumes a pre-excluded nlist and @@ -948,7 +1028,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, std::map output_map; extract_outputs(output_map, flat_outputs); - if (lower_input_is_graph_) { + if (lower_input_is_graph_ || lower_input_is_canonical_) { // 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. @@ -1097,7 +1177,16 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, "dp convert-backend --atomic-virial INPUT.pth OUTPUT.pt2"); } int natoms = atype.size(); - int nframes = coord.size() / (natoms * 3); + if (natoms == 0) { + throw deepmd::deepmd_exception( + "The coordinate-based inference interface requires at least one atom."); + } + const std::size_t frame_coord_size = static_cast(natoms) * 3; + if (coord.empty() || coord.size() % frame_coord_size != 0) { + throw deepmd::deepmd_exception( + "coord size must equal nframes * natoms * 3 with nframes >= 1."); + } + int nframes = static_cast(coord.size() / frame_coord_size); if (nframes > 1) { // Multi-frame: loop over frames and concatenate compute_nframes(ener, force, virial, atom_energy, atom_virial, nframes, @@ -1203,7 +1292,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, if (lower_input_is_edge_) { edge_tensors = createEdgeTensors(nlist_raw, coord_cpy_d, mapping_64, nloc, nall, device); - } else if (lower_input_is_graph_) { + } else if (lower_input_is_graph_ || lower_input_is_canonical_) { // Standalone (no nlist) graph schema: build_nlist already cut at rcut and // keys row i to center i, so no row_centers remapping is needed. graph_tensors = @@ -1241,18 +1330,8 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, fparam_tensor = torch::zeros({0}, options).to(device); } - at::Tensor aparam_tensor; - if (!aparam.empty()) { - aparam_tensor = - torch::from_blob( - const_cast(aparam.data()), - {1, natoms, static_cast(aparam.size()) / natoms}, - valuetype_options) - .to(torch::kFloat64) - .to(device); - } else { - aparam_tensor = torch::zeros({0}, options).to(device); - } + 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. @@ -1295,18 +1374,29 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, edge_tensors.edge_index, edge_tensors.edge_vec, edge_tensors.edge_index_ext, edge_tensors.edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor); - } else if (lower_input_is_graph_) { - // Model-level pair exclusion is a BUILD-time transform (decision - // #18/A4): the exported graph lower consumes a pre-excluded edge_mask and - // never re-applies it; this is the single application site on the C++ - // graph route. - const at::Tensor graph_edge_mask = deepmd::applyPairExclusion( + } else if (lower_input_is_graph_ || lower_input_is_canonical_) { + graph_tensors.edge_mask = deepmd::applyPairExclusion( graph_tensors.edge_index, graph_tensors.edge_mask, graph_tensors.atype, pair_exclude_table_, ntypes); - flat_outputs = run_model_graph( - graph_tensors.atype, graph_tensors.n_node, graph_tensors.edge_index, - graph_tensors.edge_vec, graph_edge_mask, fparam_tensor, aparam_tensor, - charge_spin_tensor); + canonicalizeGraphPayload(graph_tensors, graph_tensors.atype.size(0)); + if (graph_edge_fp32_) { + graph_tensors.edge_vec = graph_tensors.edge_vec.to(torch::kFloat32); + } + if (lower_input_is_canonical_) { + const auto compact = compactCanonicalGraph(graph_tensors); + flat_outputs = run_model_canonical_graph( + compact.atype, compact.n_node, compact.n_local, compact.source, + compact.edge_vec, compact.destination_row_ptr, compact.source_row_ptr, + compact.source_order); + } else { + flat_outputs = run_model_graph( + graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, + graph_tensors.edge_index, graph_tensors.edge_vec, + graph_tensors.edge_mask, graph_tensors.destination_order, + graph_tensors.destination_row_ptr, graph_tensors.source_order, + graph_tensors.source_row_ptr, fparam_tensor, aparam_tensor, + charge_spin_tensor); + } } else { // Model-level pair exclusion is a BUILD-time transform (decision // #18/A4): the exported dense lower consumes a pre-excluded nlist and @@ -1323,7 +1413,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, std::map output_map; extract_outputs(output_map, flat_outputs); - if (lower_input_is_graph_) { + if (lower_input_is_graph_ || lower_input_is_canonical_) { // The graph forward emits LOCAL public keys; rewrite them into the dense // internal-key layout used below. nloc == N (graph node count); pad the // per-atom force/virial up to the extended nall with zero ghost rows so the @@ -1395,8 +1485,15 @@ void DeepPotPTExpt::compute_nframes(ENERGYVTYPE& ener, const std::vector& charge_spin, const bool atomic) { int natoms = atype.size(); - int dap = aparam.empty() ? 0 : static_cast(aparam.size()) / nframes; - int dfp = fparam.empty() ? 0 : static_cast(fparam.size()) / nframes; + if (nframes <= 0 || natoms <= 0) { + throw deepmd::deepmd_exception( + "Multi-frame inference requires nframes >= 1 and natoms >= 1."); + } + const FrameParameterLayout aparam_layout = resolve_frame_parameter_layout( + "aparam", aparam.size(), nframes, + 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. @@ -1430,16 +1527,18 @@ void DeepPotPTExpt::compute_nframes(ENERGYVTYPE& ener, frame_box.assign(box.begin() + s_ff * 9, box.begin() + (s_ff + 1) * 9); } std::vector frame_fparam; - if (!fparam.empty()) { - size_t s_dfp = static_cast(dfp); - frame_fparam.assign(fparam.begin() + s_ff * s_dfp, - fparam.begin() + (s_ff + 1) * s_dfp); + if (fparam_layout.frame_size > 0) { + const std::size_t offset = + fparam_layout.broadcast ? 0 : s_ff * fparam_layout.frame_size; + frame_fparam.assign(fparam.begin() + offset, + fparam.begin() + offset + fparam_layout.frame_size); } std::vector frame_aparam; - if (!aparam.empty()) { - size_t s_dap = static_cast(dap); - frame_aparam.assign(aparam.begin() + s_ff * s_dap, - aparam.begin() + (s_ff + 1) * s_dap); + if (aparam_layout.frame_size > 0) { + const std::size_t offset = + aparam_layout.broadcast ? 0 : s_ff * aparam_layout.frame_size; + frame_aparam.assign(aparam.begin() + offset, + aparam.begin() + offset + aparam_layout.frame_size); } std::vector frame_chg_spin; if (!charge_spin.empty()) { @@ -1592,9 +1691,23 @@ void DeepPotPTExpt::compute_mixed_type_impl( const bool atomic) { // Mixed-type: atype has nframes * natoms elements. // Loop over frames, each with its own atype slice. + if (nframes <= 0 || atype.empty() || + atype.size() % static_cast(nframes) != 0) { + throw deepmd::deepmd_exception( + "Mixed-type inference requires nframes >= 1, natoms >= 1, and " + "atype.size() == nframes * natoms."); + } int natoms = static_cast(atype.size()) / nframes; - int dap = aparam.empty() ? 0 : static_cast(aparam.size()) / nframes; - int dfp = fparam.empty() ? 0 : static_cast(fparam.size()) / nframes; + if (coord.size() != static_cast(nframes) * + static_cast(natoms) * 3) { + throw deepmd::deepmd_exception( + "coord size must equal nframes * natoms * 3."); + } + const FrameParameterLayout aparam_layout = resolve_frame_parameter_layout( + "aparam", aparam.size(), nframes, + 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. @@ -1630,16 +1743,18 @@ void DeepPotPTExpt::compute_mixed_type_impl( frame_box.assign(box.begin() + s_ff * 9, box.begin() + (s_ff + 1) * 9); } std::vector frame_fparam; - if (!fparam.empty()) { - size_t s_dfp = static_cast(dfp); - frame_fparam.assign(fparam.begin() + s_ff * s_dfp, - fparam.begin() + (s_ff + 1) * s_dfp); + if (fparam_layout.frame_size > 0) { + const std::size_t offset = + fparam_layout.broadcast ? 0 : s_ff * fparam_layout.frame_size; + frame_fparam.assign(fparam.begin() + offset, + fparam.begin() + offset + fparam_layout.frame_size); } std::vector frame_aparam; - if (!aparam.empty()) { - size_t s_dap = static_cast(dap); - frame_aparam.assign(aparam.begin() + s_ff * s_dap, - aparam.begin() + (s_ff + 1) * s_dap); + if (aparam_layout.frame_size > 0) { + const std::size_t offset = + aparam_layout.broadcast ? 0 : s_ff * aparam_layout.frame_size; + frame_aparam.assign(aparam.begin() + offset, + aparam.begin() + offset + aparam_layout.frame_size); } std::vector frame_chg_spin; if (!charge_spin.empty()) { @@ -1818,28 +1933,67 @@ void DeepPotPTExpt::computew_mixed_type(std::vector& ener, }); } -void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, - double* d_force, - double* d_atom_virial, - const double* d_coord, - const int* d_atype, - const int* d_edge_index, - const double* d_edge_vec, - const int nloc, - const int nedge) { - // Fully device-resident edge inference for single-domain SeZM/DPA4 models. +template +void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const EDGE_TYPE* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist) { + if (lower_input_is_canonical_) { + throw deepmd::deepmd_exception( + "compute_edges_gpu cannot serve a compact canonical artifact; use " + "compute_canonical_graph_gpu."); + } + constexpr bool edge_fp32 = std::is_same_v; + if (lower_input_is_graph_ && edge_fp32 != graph_edge_fp32_) { + throw deepmd::deepmd_exception( + "compute_edges_gpu: edge-vector pointer precision does not match the " + "graph_edge_dtype recorded by the .pt2 artifact."); + } + // ``nall_nodes`` is the graph node count: ``nloc`` folds ghosts onto local + // owners (single domain, minimum image); ``nall_nodes > nloc`` keeps the + // extended (local + ghost) node set so a domain-decomposed run operates on + // per-extended forces. A non-message-passing graph then reverse-communicates + // ghost forces to their owners (the caller's responsibility); a + // message-passing model instead exchanges ghost features across ranks inside + // the forward pass, which requires ``comm_nlist`` and the with-comm artifact. + const int nnode = (nall_nodes > 0) ? nall_nodes : nloc; + const bool extended = nnode != nloc; + const bool use_with_comm = has_comm_artifact_ && comm_nlist; + if (lower_input_is_graph_ && extended && has_message_passing_) { + throw deepmd::deepmd_exception( + "compute_edges_gpu: message-passing graph-input inference requires an " + "edge-input model with its with-comm artifact."); + } + if (extended && !lower_input_is_graph_ && !use_with_comm) { + throw deepmd::deepmd_exception( + "compute_edges_gpu: the extended node set (nall_nodes > nloc) requires " + "a graph-input model, or a message-passing model with its with-comm " + "artifact and a comm_nlist."); + } + // Fully device-resident inference for edge-input and graph-input models. // - // The caller (an MD engine such as GPUMD) builds the neighbor list and the - // compact edge schema on the GPU and passes raw device pointers. This entry - // keeps every tensor on the GPU: coordinates, the edge graph and the model - // outputs never leave the device, eliminating the host neighbor-list build - // and the per-step host-device transfers of the standalone ``compute`` path. + // The caller (an MD engine or the LAMMPS Kokkos pair) builds the neighbor + // list and the compact edge schema on the GPU and passes raw device + // pointers. This entry keeps every tensor on the GPU: coordinates, the edge + // graph and the model outputs never leave the device, eliminating the host + // neighbor-list build and the per-step host-device transfers of the + // standalone ``compute`` path. // - // Edge contract (single domain, minimum-image): ``edge_index`` and - // ``edge_scatter_index`` coincide and index local atoms; ``edge_vec`` carries - // the minimum-image bond vector ``r_neighbor - r_center``. The SeZM force is - // ``dE/d(edge_vec)`` scattered through ``edge_scatter_index``, so the - // per-atom force lands directly on local atoms with no ghost fold-back. + // Edge contract: ``edge_index`` is the flattened [2, nedge] graph (row 0 = + // neighbor/source, row 1 = center/destination) and ``edge_vec`` carries the + // bond vector ``r_neighbor - r_center``. On a single domain the graph folds + // ghosts onto local owners and indexes local atoms; under domain + // decomposition it indexes the extended node set and the per-node forces are + // folded onto their owners by the caller's reverse communication. if (!gpu_enabled) { throw deepmd::deepmd_exception( "compute_edges_gpu requires a CUDA device but the model was loaded on " @@ -1852,14 +2006,20 @@ void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, } if (!lower_input_is_edge_ && !lower_input_is_graph_) { throw deepmd::deepmd_exception( - "compute_edges_gpu requires an edge-input (SeZM/DPA4) or graph-input " - "(DPA1/DPA2/DPA3) .pt2 model."); + "compute_edges_gpu requires an edge-input or graph-input .pt2 model."); + } + if (lower_input_is_graph_ && nnode == 0) { + return; } translate_error([&] { const torch::Device device(torch::kCUDA, gpu_id); const c10::DeviceGuard device_guard(device); const auto opt_f64 = torch::TensorOptions().dtype(torch::kFloat64).device(device); + const auto opt_edge = + torch::TensorOptions() + .dtype(edge_fp32 ? torch::kFloat32 : torch::kFloat64) + .device(device); const auto opt_i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); const auto opt_i64 = @@ -1868,51 +2028,70 @@ void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, torch::TensorOptions().dtype(torch::kBool).device(device); // === Step 1. Wrap caller GPU buffers as tensors (no copy) === - at::Tensor coord_t = - torch::from_blob(const_cast(d_coord), {1, nloc, 3}, opt_f64); - at::Tensor atype_t = - torch::from_blob(const_cast(d_atype), {1, nloc}, opt_i32) - .to(torch::kInt64); + at::Tensor coord_t = nloc > 0 + ? torch::from_blob(const_cast(d_coord), + {1, nloc, 3}, opt_f64) + : torch::empty({1, 0, 3}, opt_f64); + at::Tensor atype_t = nnode > 0 ? torch::from_blob(const_cast(d_atype), + {1, nnode}, opt_i32) + .to(torch::kInt64) + : torch::empty({1, 0}, opt_i64); at::Tensor edge_index_real = - torch::from_blob(const_cast(d_edge_index), {2, nedge}, opt_i32) - .to(torch::kInt64); + nedge > 0 ? torch::from_blob(const_cast(d_edge_index), {2, nedge}, + opt_i32) + .to(torch::kInt64) + : torch::empty({2, 0}, opt_i64); at::Tensor edge_vec_real = - torch::from_blob(const_cast(d_edge_vec), {nedge, 3}, opt_f64); + nedge > 0 ? torch::from_blob(const_cast(d_edge_vec), + {nedge, 3}, opt_edge) + : torch::empty({0, 3}, opt_edge); // === Step 2. Append two masked dummy edges (exported-graph contract) === at::Tensor edge_index = torch::cat({edge_index_real, torch::zeros({2, 2}, opt_i64)}, 1); at::Tensor edge_vec = - torch::cat({edge_vec_real, torch::zeros({2, 3}, opt_f64)}, 0); + torch::cat({edge_vec_real, torch::zeros({2, 3}, opt_edge)}, 0); at::Tensor edge_mask = torch::cat( {torch::ones({nedge}, opt_bool), torch::zeros({2}, opt_bool)}); // Single-domain scheme: the force-scatter index is the (local) edge graph. const at::Tensor& edge_scatter_index = edge_index; - // === Step 3. Optional model inputs (defaults for fparam/charge_spin) === + // === Step 3. Optional model inputs (runtime fparam/aparam or defaults) === + // A non-empty runtime ``fparam`` overrides the stored default; an empty one + // falls back to the ``.pt2`` metadata default. ``aparam`` is per-atom and + // has no metadata default, so it must contain exactly ``nloc * daparam`` + // values. An empty owned partition therefore carries a valid empty tensor + // with shape ``(1, 0, daparam)``. + const auto opt_f64_cpu = torch::TensorOptions().dtype(torch::kFloat64); at::Tensor fparam_tensor; if (dfparam > 0) { - if (!(has_default_fparam_ && !default_fparam_.empty())) { + const std::vector* fp = nullptr; + if (!fparam.empty()) { + if (static_cast(fparam.size()) != dfparam) { + throw deepmd::deepmd_exception( + "compute_edges_gpu: fparam size does not match the model's " + "dfparam."); + } + fp = &fparam; + } else if (has_default_fparam_ && !default_fparam_.empty()) { + fp = &default_fparam_; + } else { throw deepmd::deepmd_exception( - "compute_edges_gpu: model requires fparam but no default_fparam is " - "stored in the .pt2 metadata."); + "compute_edges_gpu: model requires fparam but none was provided " + "and " + "no default_fparam is stored in the .pt2 metadata."); } fparam_tensor = - torch::from_blob( - const_cast(default_fparam_.data()), - {1, static_cast(default_fparam_.size())}, - torch::TensorOptions().dtype(torch::kFloat64)) + torch::from_blob(const_cast(fp->data()), + {1, static_cast(fp->size())}, + opt_f64_cpu) .clone() .to(device); } else { fparam_tensor = torch::zeros({0}, opt_f64); } - if (daparam > 0) { - throw deepmd::deepmd_exception( - "compute_edges_gpu: aparam models are not supported by the GPU edge " - "path."); - } - at::Tensor aparam_tensor = torch::zeros({0}, opt_f64); + at::Tensor aparam_tensor = + make_aparam_tensor(aparam, nloc, daparam, device); at::Tensor charge_spin_tensor; if (dchgspin > 0) { if (default_chg_spin_.empty()) { @@ -1928,6 +2107,39 @@ void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, .to(device); } + // Communication tensors for the message-passing (with-comm) forward pass: + // the extended node set exchanges ghost features across ranks via + // border_op. The send/recv swap metadata comes straight from the caller's + // neighbor list; the tensors are tiny (nswap entries) and built on the + // host. + const int phantom_n = + use_with_comm && lower_input_is_edge_ && nloc == 0 ? 2 : 0; + std::vector> shifted_sendlist; + std::vector shifted_sendlist_ptrs; + std::vector comm_tensors; + if (use_with_comm) { + if (phantom_n > 0) { + shifted_sendlist.resize(comm_nlist->nswap); + shifted_sendlist_ptrs.resize(comm_nlist->nswap); + for (int iswap = 0; iswap < comm_nlist->nswap; ++iswap) { + shifted_sendlist[iswap].assign( + comm_nlist->sendlist[iswap], + comm_nlist->sendlist[iswap] + comm_nlist->sendnum[iswap]); + for (int& index : shifted_sendlist[iswap]) { + index += phantom_n; + } + shifted_sendlist_ptrs[iswap] = shifted_sendlist[iswap].data(); + } + comm_tensors = deepmd::ptexpt::build_comm_tensors_positional( + *comm_nlist, shifted_sendlist_ptrs.data(), comm_nlist->sendnum, + comm_nlist->recvnum, phantom_n, nnode); + } else { + comm_tensors = deepmd::ptexpt::build_comm_tensors_positional( + *comm_nlist, comm_nlist->sendlist, comm_nlist->sendnum, + comm_nlist->recvnum, nloc, nnode - nloc); + } + } + // === Step 4. Run the exported model and read the per-atom outputs === // The two lower forms share the masked edge tensors but differ in both the // input set and the output naming, so each form runs and unpacks itself; @@ -1936,35 +2148,296 @@ void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, at::Tensor ae, force_t, av; std::map out; if (lower_input_is_graph_) { - // Graph (DPA1/DPA2/DPA3 NeighborGraph): single-frame node count and a - // flat node-major atype; the model returns the high-level per-atom - // quantities. + // Graph-input: descriptor nodes include halos while n_local gates fitting + // outputs and the energy reduction to owned atoms. const at::Tensor n_node = + torch::full({1}, static_cast(nnode), opt_i64); + const at::Tensor n_local = torch::full({1}, static_cast(nloc), opt_i64); + at::Tensor graph_aparam = aparam_tensor; + if (daparam > 0 && nnode > nloc) { + graph_aparam = torch::cat( + {aparam_tensor, torch::zeros({1, nnode - nloc, daparam}, opt_f64)}, + 1); + } + GraphTensorPack graph_pack; + graph_pack.atype = atype_t.reshape({nnode}); + graph_pack.n_node = n_node; + graph_pack.n_local = n_local; + graph_pack.edge_index = edge_index; + graph_pack.edge_vec = edge_vec; + graph_pack.edge_mask = deepmd::applyPairExclusion( + edge_index, edge_mask, graph_pack.atype, pair_exclude_table_, ntypes); + if (pair_exclude_table_.defined()) { + canonicalizeGraphPayload(graph_pack, nnode); + } else { + // The device-edge API requires a destination-major physical-edge + // prefix. The Kokkos producer satisfies this contract, so the exported + // graph's destination-sorted specialization can use identity order. + using BuildGraphCSR = + std::tuple(torch::Tensor, c10::SymInt, c10::SymInt); + static const auto build_graph_csr = + c10::Dispatcher::singleton() + .findSchemaOrThrow("deepmd::build_graph_csr", "") + .typed(); + std::tie(graph_pack.destination_order, graph_pack.destination_row_ptr, + graph_pack.source_order, graph_pack.source_row_ptr) = + build_graph_csr.call(edge_index, c10::SymInt(nnode), + c10::SymInt(nedge)); + } extract_outputs( - out, run_model_graph(atype_t.reshape({nloc}), n_node, edge_index, - edge_vec, edge_mask, fparam_tensor, - aparam_tensor, charge_spin_tensor)); - ae = out["atom_energy"].reshape({nloc}).contiguous(); - force_t = out["force"].reshape({nloc, 3}).contiguous(); - av = out["atom_virial"].reshape({nloc, 9}).contiguous(); + out, + run_model_graph( + graph_pack.atype, graph_pack.n_node, graph_pack.n_local, + graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask, + graph_pack.destination_order, graph_pack.destination_row_ptr, + graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, + graph_aparam, charge_spin_tensor)); + ae = out["atom_energy"].reshape({nnode}).slice(0, 0, nloc).contiguous(); + force_t = out["force"].reshape({nnode, 3}).contiguous(); + av = out["atom_virial"].reshape({nnode, 9}).contiguous(); + } else if (use_with_comm && phantom_n > 0) { + // Two masked local atoms satisfy the exported lower bounds while every + // rank participates in border communication. The prefix also reserves + // the slots before received halo features. + at::Tensor coord_ext = + nnode > 0 ? torch::from_blob(const_cast(d_coord), + {1, nnode, 3}, opt_f64) + : torch::empty({1, 0, 3}, opt_f64); + at::Tensor padded_coord = + torch::cat({torch::zeros({1, phantom_n, 3}, opt_f64), coord_ext}, 1); + at::Tensor padded_atype = + torch::cat({torch::zeros({1, phantom_n}, opt_i64), atype_t}, 1); + at::Tensor local_atype = torch::zeros({1, phantom_n}, opt_i64); + at::Tensor phantom_edge_index = torch::zeros({2, 2}, opt_i64); + at::Tensor phantom_edge_vec = torch::zeros({2, 3}, opt_edge); + at::Tensor phantom_edge_mask = torch::zeros({2}, opt_bool); + at::Tensor phantom_aparam = + daparam > 0 ? torch::zeros({1, phantom_n, daparam}, opt_f64) + : aparam_tensor; + extract_outputs( + out, + run_model_edges_with_comm( + padded_coord, local_atype, padded_atype, phantom_edge_index, + phantom_edge_vec, phantom_edge_index, phantom_edge_mask, + fparam_tensor, phantom_aparam, charge_spin_tensor, comm_tensors)); + ae = out["energy"].reshape({phantom_n}).slice(0, phantom_n).contiguous(); + force_t = out["energy_derv_r"] + .squeeze(-2) + .reshape({phantom_n + nnode, 3}) + .slice(0, phantom_n) + .contiguous(); + av = out["energy_derv_c"] + .squeeze(-2) + .reshape({phantom_n + nnode, 9}) + .slice(0, phantom_n) + .contiguous(); + } else if (use_with_comm) { + // Message-passing edge model, domain-decomposed: the extended node set + // supplies coordinates and types, the local slice feeds the fitting, and + // ghost node features are exchanged across ranks inside the forward pass + // via the with-comm artifact. ``edge_index`` already indexes the extended + // set, so it doubles as the force-scatter index. + at::Tensor coord_ext = torch::from_blob(const_cast(d_coord), + {1, nnode, 3}, opt_f64); + extract_outputs( + out, run_model_edges_with_comm( + coord_ext, atype_t.slice(1, 0, nloc), atype_t, edge_index, + edge_vec, edge_index, edge_mask, fparam_tensor, + aparam_tensor, charge_spin_tensor, comm_tensors)); + ae = out["energy"].reshape({nloc}).contiguous(); + force_t = + out["energy_derv_r"].squeeze(-2).reshape({nnode, 3}).contiguous(); + av = out["energy_derv_c"].squeeze(-2).reshape({nnode, 9}).contiguous(); } else { - // Edge (SeZM/DPA4): coord + edge_scatter_index; the model returns the raw - // reduced-energy derivatives (force/virial per extended atom). + // Edge, single domain (minimum image): coordinates and edge-scatter index + // over local atoms; the model returns the reduced-energy derivatives. extract_outputs( out, run_model_edges(coord_t, atype_t, edge_index, edge_vec, edge_scatter_index, edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor)); ae = out["energy"].reshape({nloc}).contiguous(); force_t = - out["energy_derv_r"].squeeze(-2).reshape({nloc, 3}).contiguous(); - av = out["energy_derv_c"].squeeze(-2).reshape({nloc, 9}).contiguous(); + out["energy_derv_r"].squeeze(-2).reshape({nnode, 3}).contiguous(); + av = out["energy_derv_c"].squeeze(-2).reshape({nnode, 9}).contiguous(); } - // === Step 5. Copy per-atom outputs into caller GPU buffers (D2D) === - torch::from_blob(d_atom_energy, {nloc}, opt_f64).copy_(ae); - torch::from_blob(d_force, {nloc, 3}, opt_f64).copy_(force_t); - torch::from_blob(d_atom_virial, {nloc, 9}, opt_f64).copy_(av); + // === Step 5. Copy outputs into caller GPU buffers (D2D) === + // Energy is per local atom (nloc); force and per-atom virial span the node + // set (nnode == nloc folded, or nall extended). + if (nloc > 0) { + torch::from_blob(d_atom_energy, {nloc}, opt_f64).copy_(ae); + } + if (nnode > 0) { + torch::from_blob(d_force, {nnode, 3}, opt_f64).copy_(force_t); + torch::from_blob(d_atom_virial, {nnode, 9}, opt_f64).copy_(av); + } + synchronize_current_accelerator_stream(); + }); +} + +void DeepPotPTExpt::compute_canonical_graph_gpu_impl( + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::int64_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 int nloc, + const int nall_nodes, + const std::int64_t edge_storage) { + if (!lower_input_is_canonical_) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu requires a compact canonical artifact."); + } + 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) { + 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_i64 = + torch::TensorOptions().dtype(torch::kInt64).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 edge_vec = torch::from_blob(const_cast(d_edge_vec), + {edge_storage, 3}, opt_f32); + auto destination_row_ptr = + torch::from_blob(const_cast(d_destination_row_ptr), + {nall_nodes + 1}, opt_i64); + 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); + auto n_node = torch::full({1}, nall_nodes, opt_i64); + auto n_local = torch::full({1}, nloc, opt_i64); + + std::map output; + extract_outputs(output, + run_model_canonical_graph(atype, n_node, n_local, source, + edge_vec, destination_row_ptr, + source_row_ptr, source_order)); + auto atom_energy = output["atom_energy"] + .reshape({nall_nodes}) + .slice(0, 0, nloc) + .contiguous(); + auto force = output["force"].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}, + torch::TensorOptions().dtype(torch::kFloat64).device(device)) + .copy_(atom_energy); + } + torch::from_blob( + d_force, {nall_nodes, 3}, + torch::TensorOptions().dtype(torch::kFloat64).device(device)) + .copy_(force); + torch::from_blob( + d_atom_virial, {nall_nodes, 9}, + torch::TensorOptions().dtype(torch::kFloat64).device(device)) + .copy_(atom_virial); + synchronize_current_accelerator_stream(); }); } + +void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + const int nloc, + const int nedge) { + compute_edges_gpu_impl(d_atom_energy, d_force, d_atom_virial, d_coord, + d_atype, d_edge_index, d_edge_vec, nloc, nedge, + /*fparam=*/{}, /*aparam=*/{}, /*nall_nodes=*/0, + /*comm_nlist=*/nullptr); +} + +void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const double* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist) { + compute_edges_gpu_impl(d_atom_energy, d_force, d_atom_virial, d_coord, + d_atype, d_edge_index, d_edge_vec, nloc, nedge, fparam, + aparam, nall_nodes, comm_nlist); +} + +void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const double* d_coord, + const int* d_atype, + const int* d_edge_index, + const float* d_edge_vec, + const int nloc, + const int nedge, + const std::vector& fparam, + const std::vector& aparam, + const int nall_nodes, + const InputNlist* comm_nlist) { + compute_edges_gpu_impl(d_atom_energy, d_force, d_atom_virial, d_coord, + d_atype, d_edge_index, d_edge_vec, nloc, nedge, fparam, + aparam, nall_nodes, comm_nlist); +} + +void DeepPotPTExpt::compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::int64_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 int nloc, + const int nall_nodes, + const std::int64_t edge_storage) { + compute_canonical_graph_gpu_impl( + d_atom_energy, d_force, d_atom_virial, d_atype, d_source, d_edge_vec, + d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, + edge_storage); +} + +bool DeepPotPTExpt::uses_fp32_edge_vectors() const { + return (lower_input_is_graph_ || lower_input_is_canonical_) && + graph_edge_fp32_; +} + +bool DeepPotPTExpt::supports_device_edge_inference() const { + return lower_input_is_edge_ || lower_input_is_graph_ || + lower_input_is_canonical_; +} + +bool DeepPotPTExpt::uses_canonical_graph_inference() const { + return lower_input_is_canonical_; +} + #endif diff --git a/source/api_cc/src/common.cc b/source/api_cc/src/common.cc index a99f622724..e3f718e150 100644 --- a/source/api_cc/src/common.cc +++ b/source/api_cc/src/common.cc @@ -190,7 +190,7 @@ void deepmd::select_real_atoms_coord(std::vector& dcoord, // aparam if (daparam > 0) { aparam.resize(static_cast(nframes) * - (aparam_nall ? nall_real : nloc_real)); + (aparam_nall ? nall_real : nloc_real) * daparam); select_map(aparam, aparam_, fwd_map, daparam, nframes, (aparam_nall ? nall_real : nloc_real), (aparam_nall ? nall : (nall - nghost))); diff --git a/source/api_cc/src/commonPTExpt.h b/source/api_cc/src/commonPTExpt.h index 516ccce964..dfd6057c59 100644 --- a/source/api_cc/src/commonPTExpt.h +++ b/source/api_cc/src/commonPTExpt.h @@ -1,8 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later // Shared utilities for pt_expt (.pt2 / AOTInductor) backend classes. // Provides: JSON parser, ZIP archive reader, type-sorted nlist builder, -// and helpers for the with-comm dual-artifact layout (Phase 4 of the -// GNN MPI plumbing). +// and helpers for the with-comm dual-artifact layout. #pragma once #include @@ -470,7 +469,7 @@ inline std::string read_zip_entry(const std::string& zip_path, } // ============================================================================ -// With-comm artifact extraction (Phase 4) +// With-comm artifact extraction // // GNN .pt2 archives carry a nested ``extra/forward_lower_with_comm.pt2`` // alongside the regular forward_lower artifact. AOTInductor's @@ -559,7 +558,7 @@ class TempFile { }; // ============================================================================ -// comm_dict tensor packing for the with-comm artifact (Phase 4) +// Communication tensor packing for the with-comm artifact // // The with-comm AOTInductor artifact accepts comm tensors as 8 additional // positional inputs (after the regular 4-6 inputs) in this canonical order: @@ -576,9 +575,7 @@ class TempFile { // ============================================================================ /** - * @brief Build the 8 comm-tensor positional inputs from LAMMPS data - * (Phase 5 working signature, restored after the consolidation - * attempt regressed). + * @brief Build the 8 comm-tensor positional inputs from LAMMPS data. */ inline std::vector build_comm_tensors_positional( const InputNlist& lmp_list, diff --git a/source/api_cc/tests/test_deeppot_a_fparam_aparam_nframes_ptexpt.cc b/source/api_cc/tests/test_deeppot_a_fparam_aparam_nframes_ptexpt.cc index ac5fec08a4..a94105518c 100644 --- a/source/api_cc/tests/test_deeppot_a_fparam_aparam_nframes_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_a_fparam_aparam_nframes_ptexpt.cc @@ -170,3 +170,48 @@ TYPED_TEST(TestInferDeepPotAFparamAparamNFramesPtExpt, cpu_build_nlist_atomic) { EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); } } + +TYPED_TEST(TestInferDeepPotAFparamAparamNFramesPtExpt, + rejects_trailing_aparam_value) { + using VALUETYPE = TypeParam; + std::vector invalid_aparam = this->aparam; + invalid_aparam.push_back(static_cast(0)); + std::vector ener; + std::vector force; + std::vector virial; + + EXPECT_THROW(this->dp.compute(ener, force, virial, this->coord, this->atype, + this->box, this->fparam, invalid_aparam), + deepmd::deepmd_exception); +} + +TYPED_TEST(TestInferDeepPotAFparamAparamNFramesPtExpt, + broadcasts_single_frame_aparam) { + using VALUETYPE = TypeParam; + std::vector broadcast_aparam(this->aparam.begin(), + this->aparam.begin() + this->natoms); + std::vector ener; + std::vector force; + std::vector virial; + + this->dp.compute(ener, force, virial, this->coord, this->atype, this->box, + this->fparam, broadcast_aparam); + + ASSERT_EQ(ener.size(), static_cast(this->nframes)); + EXPECT_NEAR(ener[0], ener[1], EPSILON); +} + +TYPED_TEST(TestInferDeepPotAFparamAparamNFramesPtExpt, + broadcasts_single_frame_fparam) { + using VALUETYPE = TypeParam; + std::vector broadcast_fparam = {this->fparam.front()}; + std::vector ener; + std::vector force; + std::vector virial; + + this->dp.compute(ener, force, virial, this->coord, this->atype, this->box, + broadcast_fparam, this->aparam); + + ASSERT_EQ(ener.size(), static_cast(this->nframes)); + EXPECT_NEAR(ener[0], ener[1], EPSILON); +} diff --git a/source/api_cc/tests/test_deeppot_dpa1_graph_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa1_graph_ptexpt.cc index c57abaf0a0..1a22e80038 100644 --- a/source/api_cc/tests/test_deeppot_dpa1_graph_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_dpa1_graph_ptexpt.cc @@ -1,10 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later // Test C++ inference for the NeighborGraph (graph-schema) .pt2 path of the // pt_expt backend. The graph model is a dpa1(attn_layer=0) descriptor exported -// with lower_kind="graph" (gen_dpa1.py section B); this is the FIRST runtime -// exercise of the C++ graph ingestion added in PR-B Phase B2 -// (lower_input_is_graph_ / run_model_graph / buildGraphTensors / the -// compute_inner graph branch). +// with lower_kind="graph" (gen_dpa1.py section B). // // Reference values (deeppot_dpa1_graph.expected) come from an INDEPENDENT // nlist (dense-quartet) evaluation of the same weights, so a match validates @@ -101,9 +98,7 @@ deepmd::DeepPot TestInferDpa1GraphPtExpt::dp_ref; TYPED_TEST_SUITE(TestInferDpa1GraphPtExpt, ValueTypes); -// Case 1: DeepPot builds its own neighbor list and runs the standalone graph -// branch (lower_input_is_graph_, build_nlist -> buildGraphTensors). Validates -// the graph AOTI ABI/geometry against the independent nlist reference. +// DeepPot builds its own neighbor list and runs the standalone graph branch. TYPED_TEST(TestInferDpa1GraphPtExpt, cpu_build_nlist) { using VALUETYPE = TypeParam; std::vector& coord = this->coord; @@ -114,6 +109,11 @@ TYPED_TEST(TestInferDpa1GraphPtExpt, cpu_build_nlist) { double& expected_tot_e = this->expected_tot_e; std::vector& expected_tot_v = this->expected_tot_v; deepmd::DeepPot& dp = this->dp; + deepmd::DeepPot& dp_ref = this->dp_ref; + EXPECT_TRUE(dp.supports_device_edge_inference()); + EXPECT_FALSE(dp_ref.supports_device_edge_inference()); + EXPECT_FALSE(dp.uses_fp32_edge_vectors()); + EXPECT_FALSE(dp.uses_canonical_graph_inference()); double ener; std::vector force, virial; dp.compute(ener, force, virial, coord, atype, box); @@ -130,10 +130,33 @@ TYPED_TEST(TestInferDpa1GraphPtExpt, cpu_build_nlist) { } } -// Case 2: a SECOND, larger system (12 atoms, different edge count) through the -// SAME loaded graph model — proves the dynamic edge axis works in C++. The -// graph result is cross-checked against the dense nlist .pt2 (same weights); -// at non-binding sel they must agree bit-for-bit (fp64 ~1e-10). +TYPED_TEST(TestInferDpa1GraphPtExpt, rejects_aparam_for_zero_width_model) { + using VALUETYPE = TypeParam; + std::vector ener; + std::vector force; + std::vector virial; + std::vector fparam; + std::vector invalid_aparam = {static_cast(0)}; + + EXPECT_THROW(this->dp.compute(ener, force, virial, this->coord, this->atype, + this->box, fparam, invalid_aparam), + deepmd::deepmd_exception); +} + +TYPED_TEST(TestInferDpa1GraphPtExpt, rejects_empty_coordinate_system) { + using VALUETYPE = TypeParam; + double ener = 0.0; + std::vector force; + std::vector virial; + std::vector coord; + std::vector atype; + std::vector box; + + EXPECT_THROW(this->dp.compute(ener, force, virial, coord, atype, box), + deepmd::deepmd_exception); +} + +// A larger system with a different edge count exercises the dynamic edge axis. TYPED_TEST(TestInferDpa1GraphPtExpt, cpu_build_nlist_sys2_dynamic_edges) { using VALUETYPE = TypeParam; deepmd::DeepPot& dp = this->dp; @@ -168,13 +191,8 @@ TYPED_TEST(TestInferDpa1GraphPtExpt, cpu_build_nlist_sys2_dynamic_edges) { } } -// Case 3 (CRITICAL): exercise the LAMMPS compute_inner graph branch with an -// explicit InputNlist and the `ago` cache. Calling compute twice WITHOUT -// rebuilding the nlist — first ago=0 (rebuild), then ago=1 (reuse) — must give -// identical results. This is the only case that hits compute_inner + the -// member-cached mapping_ vector; the build-nlist cases above never touch it. -// Regression guard for the OOB-on-ago>0 bug fixed by caching mapping_ as a -// member (commit 7c70db47b). +// The LAMMPS compute_inner graph branch must preserve results when ``ago`` +// reuses the cached neighbor list and atom mapping. TYPED_TEST(TestInferDpa1GraphPtExpt, lammps_nlist_ago) { using VALUETYPE = TypeParam; std::vector& coord = this->coord; @@ -221,7 +239,6 @@ TYPED_TEST(TestInferDpa1GraphPtExpt, lammps_nlist_ago) { } // ago=1: reuse the cached nlist/mapping (NO rebuild). Must match again. - // This is the path that previously read the local mapping vector OOB. ener = 0.; std::fill(force_.begin(), force_.end(), 0.0); std::fill(virial.begin(), virial.end(), 0.0); @@ -240,9 +257,47 @@ TYPED_TEST(TestInferDpa1GraphPtExpt, lammps_nlist_ago) { } } -// Case 5: exercise the DeepPot::compute ATOMIC overload on the graph .pt2. -// This is the first test to reach the ``if (atomic)`` branch inside -// remap_graph_outputs_to_dense_keys (the atom_energy/atom_virial remapping). +TYPED_TEST(TestInferDpa1GraphPtExpt, multi_rank_graph_input) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + + const int nloc = coord.size() / 3; + std::vector coord_ext; + std::vector atype_ext, mapping; + std::vector > nlist_data; + _build_nlist(nlist_data, coord_ext, atype_ext, mapping, coord, + atype, box, dp.cutoff()); + const int nall = coord_ext.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, ilist.data(), numneigh.data(), + firstneigh.data()); + convert_nlist(inlist, nlist_data); + inlist.mapping = mapping.data(); + inlist.nprocs = 2; + + double energy = 0.0; + std::vector force_ext, force, virial; + ASSERT_NO_THROW(dp.compute(energy, force_ext, virial, coord_ext, atype_ext, + box, nall - nloc, inlist, 0)); + _fold_back(force, force_ext, mapping, nloc, nall, 3); + + EXPECT_LT(fabs(energy - expected_tot_e), EPSILON); + for (int ii = 0; ii < nloc * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +// Exercise the DeepPot::compute atomic overload on the graph artifact. // The per-atom reference values are already loaded from // deeppot_dpa1_graph.expected into this->expected_e and this->expected_v by // SetUp(). @@ -286,10 +341,7 @@ TYPED_TEST(TestInferDpa1GraphPtExpt, cpu_build_nlist_atomic) { } } -// Case 4: a tiny system with no in-cutoff neighbors — only the two masked -// dummy edges survive (nedge_min=2 guard / SIGFPE-edge family). The graph -// must run cleanly, produce finite, interaction-free output (zero force/virial) -// and agree with the dense reference. +// A tiny system with no in-cutoff neighbors retains only two masked guards. TYPED_TEST(TestInferDpa1GraphPtExpt, cpu_build_nlist_tiny_no_edges) { using VALUETYPE = TypeParam; deepmd::DeepPot& dp = this->dp; diff --git a/source/api_cc/tests/test_deeppot_ptexpt.cc b/source/api_cc/tests/test_deeppot_ptexpt.cc index 530e1ed9f8..b4f12d1453 100644 --- a/source/api_cc/tests/test_deeppot_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_ptexpt.cc @@ -136,6 +136,8 @@ TYPED_TEST(TestInferDeepPotAPtExpt, cpu_build_nlist) { double& expected_tot_e = this->expected_tot_e; std::vector& expected_tot_v = this->expected_tot_v; deepmd::DeepPot& dp = this->dp; + EXPECT_FALSE(dp.supports_device_edge_inference()); + EXPECT_FALSE(dp.uses_canonical_graph_inference()); double ener; std::vector force, virial; dp.compute(ener, force, virial, coord, atype, box); @@ -644,7 +646,6 @@ TYPED_TEST(TestInferDeepPotAPtExpt, print_summary) { dp.print_summary(""); } -// Regression test for the fail-fast guard hoisted in commit c80db58d. // `deeppot_sea_no_atomic_virial.pt2` is a copy of deeppot_sea.pt2 with // the do_atomic_virial=false flag patched into its metadata.json. // Calling compute() with atomic=true on this model must throw before diff --git a/source/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc index 3ddf8dcc26..ef83eae28b 100644 --- a/source/api_cc/tests/test_neighbor_list_data.cc +++ b/source/api_cc/tests/test_neighbor_list_data.cc @@ -57,8 +57,7 @@ TEST(TestNeighborListData, MakeInlistMixedEmptyAndNonemptyRows) { } // convert_nlist(jagged) must not dereference an empty row when populating -// firstneigh. Regression test for the same `&vec[0]` UB pattern fixed in -// commit 72f95f87. +// firstneigh. TEST(TestNeighborListData, ConvertNlistEmptyRows) { std::vector> input = {{}, {}, {}}; // all rows empty std::vector ilist(input.size()), numneigh(input.size()); @@ -140,6 +139,31 @@ TEST(TestNeighborListData, RoundTripWithEmptyRows) { } #ifdef BUILD_PYTORCH +TEST(TestNeighborListData, CompactCanonicalGraphDropsMaskedGuards) { + GraphTensorPack graph; + graph.atype = torch::tensor({0}, torch::kInt64); + graph.n_node = torch::tensor({1}, torch::kInt64); + graph.n_local = torch::tensor({1}, torch::kInt64); + graph.edge_index = torch::tensor({{0, 0, 0}, {0, 0, 0}}, torch::kInt64); + graph.edge_vec = torch::tensor( + {{1.5, 0.0, 0.0}, {0.0, 0.0, 0.0}, {0.0, 0.0, 0.0}}, torch::kFloat64); + graph.edge_mask = torch::tensor({true, false, false}, torch::kBool); + graph.destination_order = torch::tensor({0, 1, 2}, torch::kInt64); + graph.destination_row_ptr = torch::tensor({0, 1}, torch::kInt64); + graph.source_order = torch::tensor({0, 1, 2}, torch::kInt64); + 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.numel(), 2); + EXPECT_EQ(compact.edge_vec.scalar_type(), torch::kFloat32); + EXPECT_EQ(compact.edge_vec.size(0), 2); + EXPECT_TRUE( + torch::equal(compact.source_order, torch::tensor({0, 1}, torch::kInt64))); + EXPECT_EQ(compact.destination_row_ptr.select(0, 1).item(), 1); + EXPECT_EQ(compact.source_row_ptr.select(0, 1).item(), 1); +} + TEST(TestEdgeTensorPack, CreateEdgeTensorsUsesRowCenters) { const torch::Device device(torch::kCPU); const std::vector> nlist = {{0}, {1}}; @@ -163,6 +187,22 @@ TEST(TestEdgeTensorPack, CreateEdgeTensorsUsesRowCenters) { EXPECT_EQ(pack.edge_index.select(0, 0).select(0, 1).item(), 1); EXPECT_EQ(pack.edge_index.select(0, 1).select(0, 1).item(), 0); EXPECT_DOUBLE_EQ(pack.edge_vec.select(0, 1).select(0, 0).item(), 1.0); + + GraphTensorPack graph; + graph.edge_index = pack.edge_index; + graph.edge_vec = pack.edge_vec; + graph.edge_mask = pack.edge_mask; + buildGraphCSR(graph, 3); + EXPECT_TRUE(torch::equal(graph.destination_order, + torch::tensor({1, 0, 2, 3}, torch::kInt64))); + EXPECT_TRUE(torch::equal(graph.destination_row_ptr, + torch::tensor({0, 1, 1, 2}, torch::kInt64))); + + canonicalizeGraphPayload(graph, 3); + EXPECT_TRUE(torch::equal(graph.destination_order, + torch::tensor({0, 1, 2, 3}, torch::kInt64))); + EXPECT_TRUE(torch::equal(graph.edge_index.select(0, 1), + torch::tensor({0, 2, 0, 0}, torch::kInt64))); } TEST(TestEdgeTensorPack, CompactFiltersSkinTopologyAndAppendsDummies) { @@ -197,7 +237,38 @@ TEST(TestEdgeTensorPack, CompactFiltersSkinTopologyAndAppendsDummies) { EXPECT_EQ(compact.edge_mask.sum().item(), 2); EXPECT_FALSE(compact.edge_mask.select(0, 2).item()); EXPECT_FALSE(compact.edge_mask.select(0, 3).item()); + + GraphTensorPack graph; + graph.edge_index = compact.edge_index; + graph.edge_mask = compact.edge_mask; + buildGraphCSR(graph, 2, /*destination_sorted=*/true); + EXPECT_TRUE(torch::equal(graph.destination_order, + torch::tensor({0, 1, 2, 3}, torch::kInt64))); + EXPECT_TRUE(torch::equal(graph.destination_row_ptr, + torch::tensor({0, 1, 2}, torch::kInt64))); + EXPECT_TRUE(torch::equal(graph.source_order, + torch::tensor({1, 0, 2, 3}, torch::kInt64))); + EXPECT_TRUE(torch::equal(graph.source_row_ptr, + torch::tensor({0, 1, 2}, torch::kInt64))); } + +TEST(TestEdgeTensorPack, CanonicalizeGraphPayloadIsStableWithinDestination) { + GraphTensorPack graph; + graph.edge_index = torch::tensor({{2, 1, 0, 0}, {1, 0, 0, 0}}, + torch::TensorOptions().dtype(torch::kInt64)); + graph.edge_vec = + torch::arange(12, torch::TensorOptions().dtype(torch::kFloat64)) + .reshape({4, 3}); + graph.edge_mask = torch::tensor({true, true, true, false}); + + canonicalizeGraphPayload(graph, 3); + + EXPECT_TRUE(torch::equal(graph.edge_index.select(0, 0), + torch::tensor({1, 0, 2, 0}, torch::kInt64))); + EXPECT_TRUE(torch::equal(graph.destination_order, + torch::tensor({0, 1, 2, 3}, torch::kInt64))); +} + #endif } // namespace deepmd diff --git a/source/api_cc/tests/test_select_real_atoms.cc b/source/api_cc/tests/test_select_real_atoms.cc new file mode 100644 index 0000000000..b693a1ba10 --- /dev/null +++ b/source/api_cc/tests/test_select_real_atoms.cc @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#include + +#include +#include + +#include "common.h" + +TEST(SelectRealAtomsCoord, PreservesWideAtomicParameters) { + constexpr int nframes = 2; + constexpr int nall = 3; + constexpr int nghost = 1; + constexpr int ntypes = 2; + constexpr int daparam = 2; + + std::vector coord(nframes * nall * 3, 0.0); + std::vector atype = {0, ntypes, 1}; + std::vector input_aparam(nframes * nall * daparam); + std::iota(input_aparam.begin(), input_aparam.end(), 0.0); + + std::vector selected_coord; + std::vector selected_atype; + std::vector selected_aparam; + std::vector forward_map; + std::vector backward_map; + int selected_nghost = 0; + int selected_nall = 0; + int selected_nloc = 0; + + deepmd::select_real_atoms_coord( + selected_coord, selected_atype, selected_aparam, selected_nghost, + forward_map, backward_map, selected_nall, selected_nloc, coord, atype, + input_aparam, nghost, ntypes, nframes, daparam, nall, true); + + EXPECT_EQ(selected_nall, 2); + EXPECT_EQ(selected_nloc, 1); + EXPECT_EQ(selected_nghost, 1); + EXPECT_EQ(selected_aparam, + (std::vector{0.0, 1.0, 4.0, 5.0, 6.0, 7.0, 10.0, 11.0})); +} diff --git a/source/lmp/pair_base.cpp b/source/lmp/pair_base.cpp index 4e76d39576..115ec70fc0 100644 --- a/source/lmp/pair_base.cpp +++ b/source/lmp/pair_base.cpp @@ -358,13 +358,11 @@ PairDeepBaseModel::PairDeepBaseModel( force_unit_cvt_factor = ener_unit_cvt_factor / dist_unit_cvt_factor; restartinfo = 1; + // Enable the compute centroid/stress/atom interface for the atomic virial. #if LAMMPS_VERSION_NUMBER >= 20201130 - centroidstressflag = - CENTROID_AVAIL; // set centroidstressflag = CENTROID_AVAIL to allow the - // use of the centroid/stress/atom. Added by Davide Tisi + centroidstressflag = CENTROID_AVAIL; #else - centroidstressflag = 2; // set centroidstressflag = 2 to allow the use of the - // centroid/stress/atom. Added by Davide Tisi + centroidstressflag = 2; #endif pppmflag = 1; respa_enable = 0; @@ -397,7 +395,8 @@ PairDeepBaseModel::PairDeepBaseModel( // set comm size needed by this Pair comm_reverse = 1; - print_summary(" "); + // The model summary is emitted by the derived constructor: the referenced + // model members are not yet constructed during base construction. } void PairDeepBaseModel::print_summary(const string pre) const { @@ -464,7 +463,8 @@ void PairDeepBaseModel::allocate() { void PairDeepBaseModel::read_restart(FILE*) { is_restart = true; } void PairDeepBaseModel::write_restart(FILE*) { - // pass + // No pair state is stored in the restart; the model is reloaded from its + // path. } void PairDeepBaseModel::init_style() { diff --git a/source/lmp/pair_base.h b/source/lmp/pair_base.h index 31ffa41c39..386a4d65a7 100644 --- a/source/lmp/pair_base.h +++ b/source/lmp/pair_base.h @@ -52,8 +52,14 @@ class PairDeepBaseModel : public Pair { double ener_unit_cvt_factor, dist_unit_cvt_factor, force_unit_cvt_factor; protected: - deepmd_compat::DeepBaseModel deep_base; - deepmd_compat::DeepBaseModelDevi deep_base_model_devi; + // Bound (not copied) to the derived class's own model members. The derived + // members are still under construction when this base constructor runs, so + // copying them would read uninitialized state; a reference only records the + // address and resolves to the fully constructed model by first use. Used for + // the model summary and, in the C++ API build, the fully device-resident + // inference entry (see pair_deepmd_kokkos). + deepmd_compat::DeepBaseModel& deep_base; + deepmd_compat::DeepBaseModelDevi& deep_base_model_devi; virtual void allocate(); double** scale; unsigned numb_models; diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index f3d38183c6..3ce4705e91 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -123,7 +123,7 @@ PairDeepMD::PairDeepMD(LAMMPS* lmp) : PairDeepBaseModel( lmp, cite_user_deepmd_package, deep_pot, deep_pot_model_devi), commdata_(nullptr) { - // Constructor body can be empty + print_summary(" "); } PairDeepMD::~PairDeepMD() { @@ -186,8 +186,8 @@ double PairDeepMD::eval_energy_with_fparam( } } - // mapping (for DPA-2/3 .pt2 GNN models that gather ghost features via - // the LAMMPS atom-map; harmless for other models). + // Owner mapping for message-passing .pt2 models that gather ghost features + // through the LAMMPS atom map; unused by other models. std::vector mapping_vec(nall, -1); if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { for (size_t ii = 0; ii < nall; ++ii) { @@ -236,6 +236,14 @@ double PairDeepMD::eval_energy_with_fparam( return scale[1][1] * dener * ener_unit_cvt_factor; } +deepmd_compat::InputNlist PairDeepMD::make_comm_nlist() { + commdata_ = (CommBrickDeepMD*)comm; + return deepmd_compat::InputNlist( + 0, nullptr, nullptr, nullptr, commdata_->nswap, commdata_->sendnum, + commdata_->recvnum, commdata_->firstrecv, commdata_->sendlist, + commdata_->sendproc, commdata_->recvproc, &world, comm->nprocs); +} + void PairDeepMD::compute(int eflag, int vflag) { if (numb_models == 0) { return; @@ -249,7 +257,7 @@ void PairDeepMD::compute(int eflag, int vflag) { "centroid/stress/atom command for 9-element atomic virial."); } bool do_ghost = true; - // dpa2 communication + // Ghost communicator used to assemble the send/recv swap metadata. commdata_ = (CommBrickDeepMD*)comm; double** x = atom->x; double** f = atom->f; @@ -297,8 +305,8 @@ void PairDeepMD::compute(int eflag, int vflag) { } } - // mapping (for DPA-2/3 .pt2 GNN models that gather ghost features via - // the LAMMPS atom-map; harmless for other models). + // Owner mapping for message-passing .pt2 models that gather ghost features + // through the LAMMPS atom map; unused by other models. std::vector mapping_vec(nall, -1); if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { for (size_t ii = 0; ii < nall; ++ii) { @@ -327,7 +335,6 @@ void PairDeepMD::compute(int eflag, int vflag) { make_fparam_from_fix(fparam); } - // int ago = numb_models > 1 ? 0 : neighbor->ago; int ago = neighbor->ago; if (numb_models > 1) { if (multi_models_no_mod_devi && @@ -381,17 +388,10 @@ void PairDeepMD::compute(int eflag, int vflag) { eatom[ii] += scale[1][1] * deatom[ii] * ener_unit_cvt_factor; } } - // Added by Davide Tisi 2020 - // interface the atomic virial computed by DeepMD - // with the one used in centroid atoms + // Map the 9-component DeePMD atomic virial onto the LAMMPS centroid + // per-atom virial (xx, yy, zz, xy, xz, yz, yx, zx, zy). if (cvflag_atom) { for (int ii = 0; ii < nall; ++ii) { - // vatom[ii][0] += 1.0 * dvatom[9*ii+0]; - // vatom[ii][1] += 1.0 * dvatom[9*ii+4]; - // vatom[ii][2] += 1.0 * dvatom[9*ii+8]; - // vatom[ii][3] += 1.0 * dvatom[9*ii+3]; - // vatom[ii][4] += 1.0 * dvatom[9*ii+6]; - // vatom[ii][5] += 1.0 * dvatom[9*ii+7]; cvatom[ii][0] += scale[1][1] * dvatom[9 * ii + 0] * ener_unit_cvt_factor; // xx cvatom[ii][1] += @@ -452,18 +452,11 @@ void PairDeepMD::compute(int eflag, int vflag) { eatom[ii] += scale[1][1] * deatom[ii] * ener_unit_cvt_factor; } } - // Added by Davide Tisi 2020 - // interface the atomic virial computed by DeepMD - // with the one used in centroid atoms + // Map the 9-component DeePMD atomic virial onto the LAMMPS centroid + // per-atom virial (xx, yy, zz, xy, xz, yz, yx, zx, zy). if (cvflag_atom) { dvatom = all_atom_virial[0]; for (int ii = 0; ii < nall; ++ii) { - // vatom[ii][0] += 1.0 * dvatom[9*ii+0]; - // vatom[ii][1] += 1.0 * dvatom[9*ii+4]; - // vatom[ii][2] += 1.0 * dvatom[9*ii+8]; - // vatom[ii][3] += 1.0 * dvatom[9*ii+3]; - // vatom[ii][4] += 1.0 * dvatom[9*ii+6]; - // vatom[ii][5] += 1.0 * dvatom[9*ii+7]; cvatom[ii][0] += scale[1][1] * dvatom[9 * ii + 0] * ener_unit_cvt_factor; // xx cvatom[ii][1] += @@ -899,7 +892,7 @@ void PairDeepMD::settings(int narg, char** arg) { << setw(18 + 1) << "max_devi_f" << setw(18 + 1) << "min_devi_f" << setw(18 + 1) << "avg_devi_f"; if (out_each) { - // at this time, we don't know how many atoms + // The atom count is not known when the header is written. fp << setw(18 + 1) << "atm_devi_f(N)"; } fp << endl; diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h index 86b2017e93..c7c40d9b18 100644 --- a/source/lmp/pair_deepmd.h +++ b/source/lmp/pair_deepmd.h @@ -54,6 +54,10 @@ class PairDeepMD : public PairDeepBaseModel { protected: deepmd_compat::DeepPot deep_pot; deepmd_compat::DeepPotModelDevi deep_pot_model_devi; + // Assemble the send/recv swap metadata (a comm-only neighbor list; its + // geometry fields are unused) for the device-resident message-passing path, + // where ghost features are exchanged across ranks inside the forward pass. + deepmd_compat::InputNlist make_comm_nlist(); private: CommBrickDeepMD* commdata_; diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp new file mode 100644 index 0000000000..9ecd68f055 --- /dev/null +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -0,0 +1,966 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#ifdef LMP_KOKKOS +#include "pair_deepmd_kokkos.h" + +#include +#include +#include +#include +#include + +#include "atom.h" +#include "atom_kokkos.h" +#include "atom_masks.h" +#include "comm.h" +#include "domain.h" +#include "error.h" +#include "force.h" +#include "kokkos.h" +#include "memory_kokkos.h" +#include "neigh_list_kokkos.h" +#include "neigh_request.h" +#include "neighbor.h" + +#ifdef KOKKOS_ENABLE_CUDA +#include +#endif + +using namespace LAMMPS_NS; + +template +PairDeepMDKokkos::PairDeepMDKokkos(LAMMPS* lmp) + : PairDeepMD(lmp), + has_null_types(false), + multi_rank(false), + nloc_model(0), + nnode_model(0), + edge_capacity(0), + edge_vec_fp32(false), + canonical_graph(false), + device_path_ok(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 | ENERGY_MASK | VIRIAL_MASK; + datamask_modify = F_MASK | ENERGY_MASK | VIRIAL_MASK; + reverse_comm_device = 1; +} + +template +PairDeepMDKokkos::~PairDeepMDKokkos() { + if (copymode) { + return; + } + memoryKK->destroy_kokkos(k_eatom, eatom); +} + +template +int PairDeepMDKokkos::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); + double** f = atom->f; + 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]; + } + return m; +} + +template +void PairDeepMDKokkos::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); + double** f = atom->f; + 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++]; + } + atomKK->modified(Host, F_MASK); +} + +template +int PairDeepMDKokkos::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( + "deepmd/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(); + const int first_i = first; + Kokkos::parallel_for( + "deepmd/kk:pack_rev", Kokkos::RangePolicy(0, n), + KOKKOS_LAMBDA(const int i) { + d_buf(3 * i + 0) = f(first_i + i, 0); + d_buf(3 * i + 1) = f(first_i + i, 1); + d_buf(3 * i + 2) = f(first_i + i, 2); + }); + return n * 3; +} + +template +void PairDeepMDKokkos::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( + "deepmd/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(); + Kokkos::parallel_for( + "deepmd/kk:unpack_rev", Kokkos::RangePolicy(0, n), + KOKKOS_LAMBDA(const int i) { + const int j = d_list(i); + f(j, 0) += d_buf(3 * i + 0); + f(j, 1) += d_buf(3 * i + 1); + f(j, 2) += d_buf(3 * i + 2); + }); +} + +template +void PairDeepMDKokkos::init_style() { + // Base setup and the full neighbor-list request. + PairDeepMD::init_style(); + + // The device edge path requires a GPU execution space and a single model. + if (std::is_same::value) { + error->all(FLERR, "pair style deepmd/kk runs on the GPU backend only."); + } + device_path_ok = deep_pot.supports_device_edge_inference(); + if (!device_path_ok) { + error->all( + FLERR, + "pair style deepmd/kk requires an edge-input or graph-input .pt2 " + "artifact; use pair style deepmd for an nlist artifact."); + } + // Domain decomposition uses the extended (local + ghost) node set: the model + // computes per-node forces and the reverse communication folds the ghost + // forces onto their owners. A single rank uses the folded minimum-image node + // set. Message-passing models additionally receive communication metadata + // through compute_edges_gpu. + multi_rank = (comm->nprocs > 1); + if (numb_models != 1) { + error->all(FLERR, "pair style deepmd/kk does not support model deviation."); + } + // A local edge graph folds ghost neighbours onto their local owner through + // the atom map; without it the fold returns -1 and corrupts the graph. + if (atom->map_style == Atom::MAP_NONE) { + error->all(FLERR, + "pair style deepmd/kk needs an atom map; add 'atom_modify map " + "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."); + } + + // 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(); + edge_vec_fp32 = deep_pot.uses_fp32_edge_vectors(); + canonical_graph = deep_pot.uses_canonical_graph_inference(); + // Force exchange transfers three values 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 = 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); + }); + } +} + +template +int PairDeepMDKokkos::build_edges_device() { + const int nlocal = atom->nlocal; + prepare_model_nodes(); + + // === Neighbor list and atom views on the device === + NeighListKokkos* 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(); + + 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; + + if ((int)d_edge_offset.extent(0) < nlocal + 1) { + d_edge_offset = Kokkos::View( + "deepmd/kk:edge_offset", nlocal + 1); + } + auto edge_offset = d_edge_offset; + + // === Pass 1: per-center edge count (0 for a virtual center) === + // A neighbour is the ghost's owner node (folded) or the ghost's own node + // (extended); a neighbour imaging a virtual atom is skipped. + Kokkos::parallel_for( + "deepmd/kk:count", Kokkos::RangePolicy(0, inum), + KOKKOS_LAMBDA(const int ii) { + const int i = d_ilist(ii); + if (loc2model(i) < 0) { + edge_offset(i) = 0; + return; + } + const double xi = x(i, 0), yi = x(i, 1), zi = x(i, 2); + const int jnum = d_numneigh(i); + int c = 0; + for (int jj = 0; jj < jnum; ++jj) { + int j = d_neighbors(i, jj); + j &= NEIGHMASK; + if (loc2model(multi ? j : owner(j)) < 0) { + continue; + } + const double dx = x(j, 0) - xi, dy = x(j, 1) - yi, dz = x(j, 2) - zi; + if (dx * dx + dy * dy + dz * dz < cutsq) { + ++c; + } + } + edge_offset(i) = c; + }); + + // === Exclusive prefix sum of the counts -> edge offsets, total = nedge === + // An empty subdomain (nlocal == 0) has no edges and must not read the scan + // sentinel, which the empty scan would leave uninitialized. + std::int64_t nedge_total = 0; + if (nlocal > 0) { + Kokkos::parallel_scan( + "deepmd/kk:scan", Kokkos::RangePolicy(0, nlocal), + KOKKOS_LAMBDA(const int i, std::int64_t& update, const bool final) { + const std::int64_t c = edge_offset(i); + if (final) { + edge_offset(i) = update; + } + update += c; + if (final && i == nlocal - 1) { + edge_offset(nlocal) = update; + } + }); + Kokkos::deep_copy(nedge_total, Kokkos::subview(d_edge_offset, nlocal)); + } + if (nedge_total > std::numeric_limits::max()) { + error->one( + FLERR, + "The DeePMD Kokkos edge count exceeds the int32 graph-index limit"); + } + const int nedge = static_cast(nedge_total); + + // Keep the edge buffers allocated (non-null) even with no physical edges, so + // an isolated-atom step still runs the model for its per-atom energy bias. + const int want = nedge > 0 ? nedge : 1; + if (edge_capacity < want) { + const std::int64_t grown = static_cast(want) + want / 8 + 64; + edge_capacity = static_cast(grown > std::numeric_limits::max() + ? std::numeric_limits::max() + : grown); + // Size in size_t: at multi-million-atom scale 3 * edge_capacity exceeds the + // 32-bit range even though the edge count itself still fits int. + d_edge_index = Kokkos::View( + "deepmd/kk:edge_index", static_cast(edge_capacity) * 2); + if (edge_vec_fp32) { + d_edge_vec_float = Kokkos::View( + "deepmd/kk:edge_vec_float", + static_cast(edge_capacity) * 3); + d_edge_vec = Kokkos::View(); + } else { + d_edge_vec = Kokkos::View( + "deepmd/kk:edge_vec", static_cast(edge_capacity) * 3); + d_edge_vec_float = Kokkos::View(); + } + } + auto edge_index = d_edge_index; + auto edge_vec = d_edge_vec; + auto edge_vec_float = d_edge_vec_float; + const bool write_fp32_edge = edge_vec_fp32; + const double inv_dist = 1.0 / dist_unit_cvt_factor; + const std::int64_t nedge_l = nedge; + + // === Pass 2: emit model-space edges (src = neighbour node, dst = center + // node; bond = x[j] - x[i]) === + Kokkos::parallel_for( + "deepmd/kk: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), yi = x(i, 1), zi = x(i, 2); + const int jnum = d_numneigh(i); + std::int64_t e = edge_offset(i); + for (int jj = 0; jj < jnum; ++jj) { + int j = d_neighbors(i, jj); + j &= NEIGHMASK; + const int mj = loc2model(multi ? j : owner(j)); + if (mj < 0) { + continue; + } + const double dx = x(j, 0) - xi, dy = x(j, 1) - yi, dz = x(j, 2) - zi; + if (dx * dx + dy * dy + dz * dz < cutsq) { + edge_index(e) = mj; // src (model node) + edge_index(nedge_l + e) = mi; // dst (model node) + if (write_fp32_edge) { + edge_vec_float(3 * e + 0) = static_cast(dx * inv_dist); + edge_vec_float(3 * e + 1) = static_cast(dy * inv_dist); + edge_vec_float(3 * e + 2) = static_cast(dz * inv_dist); + } else { + edge_vec(3 * e + 0) = dx * inv_dist; + edge_vec(3 * e + 1) = dy * inv_dist; + edge_vec(3 * e + 2) = dz * inv_dist; + } + ++e; + } + } + }); + + // 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 ((int)d_coord_model.extent(0) < 3 * nnode_model) { + d_coord_model = Kokkos::View("deepmd/kk:coord_model", + 3 * nnode_model); + } + auto coord_model = d_coord_model; + Kokkos::parallel_for( + "deepmd/kk:coord", Kokkos::RangePolicy(0, nnode_model), + KOKKOS_LAMBDA(const int m) { + const int i = model2loc(m); + coord_model(3 * m + 0) = x(i, 0); + coord_model(3 * m + 1) = x(i, 1); + coord_model(3 * m + 2) = x(i, 2); + }); + } + 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) { + error->all(FLERR, + "pair style deepmd/kk cannot execute this model input schema."); + } + 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, "deepmd/kk:eatom"); + d_eatom = k_eatom.template view(); + } + std::int64_t nedge = 0; + if (canonical_graph) { + nedge = build_canonical_edges_device(canonical_workspace); + } 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) + + // 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 + // 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("deepmd/kk:atom_energy", nlocal); + } + if ((int)d_out_force.extent(0) < 3 * nall) { + d_out_force = + Kokkos::View("deepmd/kk:out_force", 3 * nall); + d_atom_virial = + Kokkos::View("deepmd/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("deepmd/kk:reverse_virial", 9 * nall); + } + Kokkos::deep_copy(d_out_force, 0.0); + Kokkos::deep_copy(d_atom_energy, 0.0); + Kokkos::deep_copy(d_atom_virial, 0.0); + + atomKK->sync(execution_space, X_MASK | TYPE_MASK); + auto x = atomKK->k_x.template view(); + + // Runtime frame (fparam) and per-atom (aparam) parameters, built per step + // from the same sources as the standalone pair (compute / fix / ttm or a + // uniform setting). Empty vectors fall back to the model's stored defaults. + std::vector aparam_step; + if (do_compute_aparam) { + make_aparam_from_compute(aparam_step); + } else if (aparam.size() > 0) { + make_uniform_aparam(aparam_step, aparam, nlocal); + } else if (do_ttm) { +#ifdef USE_TTM + if (dim_aparam > 0) { + make_ttm_aparam(aparam_step); + } else if (dim_fparam > 0) { + make_ttm_fparam(fparam); + } +#endif + } + if (do_compute_fparam) { + make_fparam_from_compute(fparam); + } else if (do_fix_fparam) { + make_fparam_from_fix(fparam); + } + + // ``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(); + std::vector aparam_model(static_cast(nloc_m) * + dim_aparam); + for (int m = 0; m < nloc_m; ++m) { + const int i = h_m2l(m); + for (int k = 0; k < dim_aparam; ++k) { + aparam_model[static_cast(m) * dim_aparam + k] = + aparam_step[static_cast(i) * dim_aparam + k]; + } + } + aparam_step.swap(aparam_model); + } + + // Send/recv swap metadata for a message-passing model under domain + // decomposition: ghost features are exchanged across ranks inside the forward + // pass. It is passed only when the raw LAMMPS atom indices in the swap lists + // match the model-node indices, i.e. when no virtual (NULL-type) atoms + // 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) { + comm_list = make_comm_nlist(); + comm_ptr = &comm_list; + } + + if ((canonical_graph && nnode_m > 0) || 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 + // the two runtimes are bracketed by explicit synchronization: fence the + // Kokkos work before the model reads the edges, and synchronize the device + // after so the scatter sees the finished model outputs. + Kokkos::fence(); + // 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(); + 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); + } 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); + } 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); + } + } 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 DualView + // and the completed fold is copied back once after all communication stages. + auto model2loc = d_model2loc; + const double fscale = scale[1][1] * force_unit_cvt_factor; + reverse_virial = false; + reverse_used_host = false; + // The KOKKOS package runs 'newton off', so the integrator's force_clear only + // zeros the local forces (f[0, nlocal)); the ghost slots f[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); + auto f = atomKK->k_f.template view(); + auto out_force = d_out_force; + if (multi_rank) { + Kokkos::parallel_for( + "deepmd/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; + }); + } + Kokkos::parallel_for( + "deepmd/kk:scatter_f", Kokkos::RangePolicy(0, nnode_m), + KOKKOS_LAMBDA(const int m) { + const int i = model2loc(m); + 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); + }); + atomKK->modified(execution_space, F_MASK); + if (multi_rank) { + comm->reverse_comm(this, 3); + if (reverse_used_host) { + atomKK->sync(execution_space, F_MASK); + } + } + + if (eflag_global) { + auto atom_energy = d_atom_energy; + double e_sum = 0.0; + Kokkos::parallel_reduce( + "deepmd/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( + "deepmd/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( + "deepmd/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 = 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 PairDeepMDKokkos; +#ifdef LMP_KOKKOS_GPU +template class PairDeepMDKokkos; +#endif +} // namespace LAMMPS_NS + +#endif // LMP_KOKKOS diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h new file mode 100644 index 0000000000..d09e9a7abc --- /dev/null +++ b/source/lmp/pair_deepmd_kokkos.h @@ -0,0 +1,155 @@ +// 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(deepmd/kk,PairDeepMDKokkos); +PairStyle(deepmd/kk/device,PairDeepMDKokkos); +PairStyle(deepmd/kk/host,PairDeepMDKokkos); +// clang-format on +#else + +#ifndef LMP_PAIR_DEEPMD_KOKKOS_H +#define LMP_PAIR_DEEPMD_KOKKOS_H + +#include +#include + +#include "kokkos_base.h" +#include "kokkos_type.h" +#include "neigh_list_kokkos.h" +#include "pair_deepmd.h" + +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 +// ``DeepPot::compute_edges_gpu``. +// +// The neighbor list, the compact edge schema and the model outputs all stay +// on the device: the edge graph is built from the Kokkos device neighbor +// list, handed to ``compute_edges_gpu`` as raw device pointers, and the +// returned per-atom force / energy / virial are scattered back into the +// Kokkos atom arrays without any host round-trip. This removes the per-step +// host coordinate marshaling and the host-device transfers of the standalone +// ``pair_style deepmd`` 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 forces onto their owners +// through reverse communication. +template +class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { + public: + typedef DeviceType device_type; + typedef ArrayTypes AT; + + PairDeepMDKokkos(class LAMMPS*); + ~PairDeepMDKokkos() 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 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; + + // 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(); + 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; + 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; + Kokkos::View + d_coord_model; // (3 * nnode_model), NULL case + + // Compact edge schema: edge_index is [2 * nedge] (src rows then dst rows), + // edge_vec is [3 * nedge]; offsets is the per-atom exclusive edge prefix. + Kokkos::View d_edge_offset; // (nlocal + 1) + Kokkos::View d_edge_index; // (2 * nedge) + 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). + Kokkos::View d_atom_energy; // (nlocal) + Kokkos::View d_out_force; // (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; + + int edge_capacity; // allocated edges in d_edge_index / d_edge_vec + bool edge_vec_fp32; // model graph ABI consumes edge vectors in fp32 + bool canonical_graph; // compact source-only graph artifact + bool device_path_ok; // resolved once in init_style + 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/lmp/pair_deepspin.cpp b/source/lmp/pair_deepspin.cpp index e68848ab06..4de54c9197 100644 --- a/source/lmp/pair_deepspin.cpp +++ b/source/lmp/pair_deepspin.cpp @@ -120,7 +120,7 @@ static const char cite_user_deepmd_package[] = PairDeepSpin::PairDeepSpin(LAMMPS* lmp) : PairDeepBaseModel( lmp, cite_user_deepmd_package, deep_spin, deep_spin_model_devi) { - // Constructor body can be empty + print_summary(" "); } PairDeepSpin::~PairDeepSpin() { @@ -140,7 +140,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { "centroid/stress/atom command for 9-element atomic virial."); } bool do_ghost = true; - // dpa2 communication + // Ghost communicator used to assemble the send/recv swap metadata. commdata_ = (CommBrickDeepSpin*)comm; double** x = atom->x; double** f = atom->f; @@ -201,8 +201,8 @@ void PairDeepSpin::compute(int eflag, int vflag) { } } - // mapping (for DPA-2/3 .pt2 GNN models that gather ghost features via - // the LAMMPS atom-map; harmless for other models). + // Owner mapping for message-passing .pt2 models that gather ghost features + // through the LAMMPS atom map; unused by other models. std::vector mapping_vec(nall, -1); if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { for (size_t ii = 0; ii < nall; ++ii) { @@ -229,7 +229,6 @@ void PairDeepSpin::compute(int eflag, int vflag) { make_fparam_from_compute(fparam); } - // int ago = numb_models > 1 ? 0 : neighbor->ago; int ago = neighbor->ago; if (numb_models > 1) { if (multi_models_no_mod_devi && @@ -283,17 +282,10 @@ void PairDeepSpin::compute(int eflag, int vflag) { eatom[ii] += scale[1][1] * deatom[ii] * ener_unit_cvt_factor; } } - // Added by Davide Tisi 2020 - // interface the atomic virial computed by DeepMD - // with the one used in centroid atoms + // Map the 9-component DeePMD atomic virial onto the LAMMPS centroid + // per-atom virial (xx, yy, zz, xy, xz, yz, yx, zx, zy). if (cvflag_atom) { for (int ii = 0; ii < nall; ++ii) { - // vatom[ii][0] += 1.0 * dvatom[9*ii+0]; - // vatom[ii][1] += 1.0 * dvatom[9*ii+4]; - // vatom[ii][2] += 1.0 * dvatom[9*ii+8]; - // vatom[ii][3] += 1.0 * dvatom[9*ii+3]; - // vatom[ii][4] += 1.0 * dvatom[9*ii+6]; - // vatom[ii][5] += 1.0 * dvatom[9*ii+7]; cvatom[ii][0] += scale[1][1] * dvatom[9 * ii + 0] * ener_unit_cvt_factor; // xx cvatom[ii][1] += @@ -355,18 +347,11 @@ void PairDeepSpin::compute(int eflag, int vflag) { eatom[ii] += scale[1][1] * deatom[ii] * ener_unit_cvt_factor; } } - // Added by Davide Tisi 2020 - // interface the atomic virial computed by DeepMD - // with the one used in centroid atoms + // Map the 9-component DeePMD atomic virial onto the LAMMPS centroid + // per-atom virial (xx, yy, zz, xy, xz, yz, yx, zx, zy). if (cvflag_atom) { dvatom = all_atom_virial[0]; for (int ii = 0; ii < nall; ++ii) { - // vatom[ii][0] += 1.0 * dvatom[9*ii+0]; - // vatom[ii][1] += 1.0 * dvatom[9*ii+4]; - // vatom[ii][2] += 1.0 * dvatom[9*ii+8]; - // vatom[ii][3] += 1.0 * dvatom[9*ii+3]; - // vatom[ii][4] += 1.0 * dvatom[9*ii+6]; - // vatom[ii][5] += 1.0 * dvatom[9*ii+7]; cvatom[ii][0] += scale[1][1] * dvatom[9 * ii + 0] * ener_unit_cvt_factor; // xx cvatom[ii][1] += @@ -482,7 +467,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { << " " << setw(18) << all_fm_min << " " << setw(18) << all_fm_avg; } if (out_each == 1) { - // need support for spin atomic force. + // Only the per-atom force deviation is gathered here. vector std_f_all(atom->natoms); // Gather std_f and tags tagint* tag = atom->tag; diff --git a/source/lmp/plugin/CMakeLists.txt b/source/lmp/plugin/CMakeLists.txt index 8f32af3e3e..8dcbe6dd59 100644 --- a/source/lmp/plugin/CMakeLists.txt +++ b/source/lmp/plugin/CMakeLists.txt @@ -116,6 +116,15 @@ if(DEFINED LAMMPS_SOURCE_ROOT OR DEFINED LAMMPS_VERSION) add_library(${libname} MODULE ${LMP_SRC}) + option(DEEPMD_LAMMPS_KOKKOS + "Build the Kokkos device pair style in the LAMMPS plugin" OFF) + if(DEEPMD_LAMMPS_KOKKOS) + find_package(Kokkos CONFIG REQUIRED) + target_compile_definitions(${libname} PRIVATE LMP_KOKKOS) + target_include_directories(${libname} PRIVATE ${LAMMPS_HEADER_DIR}/KOKKOS) + target_link_libraries(${libname} PUBLIC Kokkos::kokkos) + endif() + # link: libdeepmd if(DP_USING_C_API) target_link_libraries(${libname} PUBLIC ${LIB_DEEPMD_C}) diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 3bb34e622d..a6560e4fec 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -1,8 +1,64 @@ file(GLOB OP_SRC print_summary.cc comm.cc tabulate_multi_device.cc) +option( + DEEPMD_CUDA_PORTABLE_PTX + "Embed a lowest-supported PTX fallback in the PyTorch CUDA operator library" + ON) +# Fused graph-lower inference operators (CUDA / cuBLAS). They include ATen CUDA +# headers and link libtorch_cuda, so they build only against a CUDA-enabled +# PyTorch (DEEPMD_TORCH_HAS_CUDA); against a CPU-only torch they are omitted and +# the Python dispatch falls back to the reference path (see +# deepmd.kernels.cuda.*.op_available). The CUDA language is scoped per +# directory, so it must be enabled here: the sibling GPU library turns it on +# only within its own subtree, which does not cover this target. +if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) + find_package(CUDAToolkit REQUIRED) + if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + # CUDA 12.9 CCCL fails to compile CUB/Thrust with -arch=all. + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.9" AND CUDAToolkit_VERSION + VERSION_LESS "13.0") + set(CMAKE_CUDA_ARCHITECTURES all-major) + else() + set(CMAKE_CUDA_ARCHITECTURES all) + endif() + endif() + enable_language(CUDA) + list( + APPEND + OP_SRC + dpa1_graph_descriptor.cu + dpa1_graph_compress.cu + graph_fitting.cu + edge_force_virial.cu + dpa1_graph_energy_force.cu) +endif() add_library(deepmd_op_pt MODULE ${OP_SRC}) # link: libdeepmd libtorch target_link_libraries(deepmd_op_pt PRIVATE ${TORCH_LIBRARIES}) +if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) + if(DEEPMD_CUDA_PORTABLE_PTX AND NOT CMAKE_CUDA_ARCHITECTURES MATCHES + "^(all|all-major)$") + if(CMAKE_CUDA_ARCHITECTURES AND NOT CMAKE_CUDA_ARCHITECTURES STREQUAL "OFF") + set(DEEPMD_PT_CUDA_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES}) + else() + set(DEEPMD_PT_CUDA_ARCHITECTURES) + endif() + if(CUDAToolkit_VERSION VERSION_LESS "13.0") + list(APPEND DEEPMD_PT_CUDA_ARCHITECTURES "70-virtual") + else() + list(APPEND DEEPMD_PT_CUDA_ARCHITECTURES "75-virtual") + endif() + list(REMOVE_DUPLICATES DEEPMD_PT_CUDA_ARCHITECTURES) + set_property(TARGET deepmd_op_pt PROPERTY CUDA_ARCHITECTURES + ${DEEPMD_PT_CUDA_ARCHITECTURES}) + endif() + target_link_libraries(deepmd_op_pt PRIVATE CUDA::cublas) + # libtorch headers require C++17; the CUDA sources must match. + set_target_properties(deepmd_op_pt PROPERTIES CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON) + set_source_files_properties(dpa1_graph_compress.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}) else() @@ -11,7 +67,8 @@ endif() remove_definitions(-D_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI}) target_compile_definitions( deepmd_op_pt - PUBLIC "$<$:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_PT}>") + PUBLIC "$<$:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_PT}>" + "$<$:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_PT}>") if(APPLE) set_target_properties(deepmd_op_pt PROPERTIES INSTALL_RPATH "@loader_path") else() diff --git a/source/op/pt/dpa1_graph_common.cuh b/source/op/pt/dpa1_graph_common.cuh new file mode 100644 index 0000000000..c8fc3f8f01 --- /dev/null +++ b/source/op/pt/dpa1_graph_common.cuh @@ -0,0 +1,564 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Shared device / host infrastructure of the DPA1 graph-lower CUDA operators. +// +// The MLP descriptor (dpa1_graph_descriptor.cu) and the tabulated compressed +// descriptor (dpa1_graph_compress.cu) evaluate the same environment matrix, +// CSR run structure, moment layout and G^T G contraction; they differ only in +// how the per-edge embedding ``g`` is produced (three-layer MLP versus quintic +// spline table). This header carries everything outside that difference: +// +// * the per-edge staging (environment-matrix row, type-pair index, smooth +// switch, center node) and the parallel CSR run scan over a tile; +// * the sort-free CSR edge ordering (histogram + exclusive scan + scatter); +// * the gram contraction kernels (forward and backward); +// * float4 fragment load / store helpers and the launch utilities. +// +// Everything lives in an anonymous namespace: each translation unit +// instantiates its own internal copy, which keeps the operators independently +// compilable (and their heavy template sets compiling in parallel). + +#pragma once + +#include +#include +#include + +#include +#include + +namespace { + +constexpr int kThreads = 256; // 8 warps per CTA +constexpr int kTileEdges = 128; // BE: edges per forward tile pass +constexpr int kTileStride = kTileEdges + 4; // padded [channel][edge] stride + +#define DEV_INLINE __device__ __forceinline__ + +#define DPA1_CHECK_LAUNCH(what) \ + do { \ + cudaError_t err = cudaGetLastError(); \ + TORCH_CHECK(err == cudaSuccess, what, ": ", cudaGetErrorString(err)); \ + } while (0) + +// Quintic smooth switch on [rmin, rmax] and its derivative. +DEV_INLINE float switch_val(float r, float rmin, float rmax) { + float u = (fminf(fmaxf(r, rmin), rmax) - rmin) / (rmax - rmin); + float u2 = u * u; + return u2 * u * (-6.f * u2 + 15.f * u - 10.f) + 1.f; +} + +DEV_INLINE float switch_deriv(float r, float rmin, float rmax) { + if (r <= rmin || r >= rmax) { + return 0.f; + } + float u = (r - rmin) / (rmax - rmin), u2 = u * u; + return (-30.f * u2 * u2 + 60.f * u2 * u - 30.f * u2) / (rmax - rmin); +} + +// Eight-row float4 fragment load / store on a 16-byte-aligned span. +DEV_INLINE void load8(const float* p, float (&a)[8]) { + const float4 v0 = *reinterpret_cast(p); + const float4 v1 = *reinterpret_cast(p + 4); + a[0] = v0.x; + a[1] = v0.y; + a[2] = v0.z; + a[3] = v0.w; + a[4] = v1.x; + a[5] = v1.y; + a[6] = v1.z; + a[7] = v1.w; +} + +DEV_INLINE void store8(float* p, const float (&a)[8]) { + float4 v0, v1; + v0.x = a[0]; + v0.y = a[1]; + v0.z = a[2]; + v0.w = a[3]; + v1.x = a[4]; + v1.y = a[5]; + v1.z = a[6]; + v1.w = a[7]; + *reinterpret_cast(p) = v0; + *reinterpret_cast(p + 4) = v1; +} + +// Streaming (evict-first) variants for the once-per-pass spill tensors; the +// cache hint keeps them from evicting the L1-resident weight lines. +DEV_INLINE void load8_streaming(const float* p, float (&a)[8]) { + const float4 v0 = __ldcs(reinterpret_cast(p)); + const float4 v1 = __ldcs(reinterpret_cast(p) + 1); + a[0] = v0.x; + a[1] = v0.y; + a[2] = v0.z; + a[3] = v0.w; + a[4] = v1.x; + a[5] = v1.y; + a[6] = v1.z; + a[7] = v1.w; +} + +DEV_INLINE void store8_streaming(float* p, const float (&a)[8]) { + float4 v0, v1; + v0.x = a[0]; + v0.y = a[1]; + v0.z = a[2]; + v0.w = a[3]; + v1.x = a[4]; + v1.y = a[5]; + v1.z = a[6]; + v1.w = a[7]; + __stcs(reinterpret_cast(p), v0); + __stcs(reinterpret_cast(p) + 1, v1); +} + +// Warp-uniform float4 weight broadcast (all lanes read the same address, so +// the load costs one L1 transaction per warp). +DEV_INLINE float4 weight4(const float* p) { + return __ldg(reinterpret_cast(p)); +} + +// Generic N-row float4 fragment load / store (N a multiple of 4). These back +// the edge-fragment loads of the backward kernels, whose per-thread edge count +// (EPT) is a template parameter; the forwards keep the fixed load8 / store8. +template +DEV_INLINE void loadN(const float* p, float (&a)[N]) { +#pragma unroll + for (int i = 0; i < N; i += 4) { + const float4 v = *reinterpret_cast(p + i); + a[i] = v.x; + a[i + 1] = v.y; + a[i + 2] = v.z; + a[i + 3] = v.w; + } +} + +template +DEV_INLINE void storeN(float* p, const float (&a)[N]) { +#pragma unroll + for (int i = 0; i < N; i += 4) { + *reinterpret_cast(p + i) = + make_float4(a[i], a[i + 1], a[i + 2], a[i + 3]); + } +} + +template +DEV_INLINE void loadN_streaming(const float* p, float (&a)[N]) { +#pragma unroll + for (int i = 0; i < N; i += 4) { + const float4 v = __ldcs(reinterpret_cast(p + i)); + a[i] = v.x; + a[i + 1] = v.y; + a[i + 2] = v.z; + a[i + 3] = v.w; + } +} + +template +DEV_INLINE void storeN_streaming(float* p, const float (&a)[N]) { +#pragma unroll + for (int i = 0; i < N; i += 4) { + __stcs(reinterpret_cast(p + i), + make_float4(a[i], a[i + 1], a[i + 2], a[i + 3])); + } +} + +__host__ __device__ constexpr long ceil_div(long a, long b) { + return (a + b - 1) / b; +} + +// ====================================================================== +// Sort-free CSR ordering: histogram + exclusive scan + scatter. +// +// Graph builders emit real edges already sorted by the center (dst) node +// with a masked padding tail, which keeps the identity permutation (masked +// rows contribute exactly zero to the moment, so their position within the +// order is irrelevant). A genuinely unsorted stream falls back to the +// scatter, which makes each node's edges contiguous. +// ====================================================================== +__global__ void edge_order_count_kernel(long n_edge, + const long* __restrict__ dst, + const bool* __restrict__ mask, + int* __restrict__ counts, + int* __restrict__ unsorted) { + const long e = blockIdx.x * (long)blockDim.x + threadIdx.x; + if (e >= n_edge) { + return; + } + const long d = dst[e]; + atomicAdd(&counts[d], 1); + if (e > 0 && mask[e] && mask[e - 1] && dst[e - 1] > d) { + atomicOr(unsorted, 1); + } +} + +__global__ void edge_order_scatter_kernel(long n_edge, + const long* __restrict__ dst, + const int* __restrict__ offsets, + const int* __restrict__ unsorted, + int* __restrict__ cursor, + int* __restrict__ order) { + const long e = blockIdx.x * (long)blockDim.x + threadIdx.x; + if (e >= n_edge) { + return; + } + if (*unsorted == 0) { + order[e] = (int)e; + return; + } + const long d = dst[e]; + const int pos = atomicAdd(&cursor[d], 1); + order[offsets[d] + pos] = (int)e; +} + +// ====================================================================== +// Per-edge staging: environment-matrix row, type-pair index, switch value +// and center node. +// ====================================================================== +struct EdgeStage { + float r0, r1, r2, r3; // normalized environment-matrix row + float sw; // raw smooth-switch value (type-pair gate factor) + int pair_idx; + int dst; + bool valid; +}; + +DEV_INLINE EdgeStage stage_edge(long e, + long n_edge, + int ntypes, + int one_side, + float rcut, + float rcut_smth, + float protection, + const float* __restrict__ edge_vec, + const long* __restrict__ edge_index, + const bool* __restrict__ edge_mask, + const long* __restrict__ atype, + const float* __restrict__ davg, + const float* __restrict__ inv_dstd) { + EdgeStage s; + const long src = edge_index[e]; + const long dst = edge_index[n_edge + e]; + s.dst = (int)dst; + const int ct = (int)atype[dst]; + const int nt = (int)atype[src]; + s.valid = edge_mask[e]; + const float x = edge_vec[e * 3 + 0]; + const float y = edge_vec[e * 3 + 1]; + const float z = edge_vec[e * 3 + 2]; + float len = sqrtf(x * x + y * y + z * z); + // Reference guard: padding edges enter the smooth switch at |r| + 1 so the + // 1/q factors stay finite; their moment weight is zeroed separately. + len += s.valid ? 0.f : 1.f; + const float q = len + protection; + const float sw = switch_val(len, rcut_smth, rcut); + const float rq = 1.f / q; + const float t0 = sw * rq, iq2 = sw * rq * rq; + const float* av = davg + (long)ct * 4; + const float* isd = inv_dstd + (long)ct * 4; + s.r0 = (t0 - av[0]) * isd[0]; + s.r1 = (x * iq2 - av[1]) * isd[1]; + s.r2 = (y * iq2 - av[2]) * isd[2]; + s.r3 = (z * iq2 - av[3]) * isd[3]; + s.sw = sw; + s.pair_idx = one_side ? nt : ct * ntypes + nt; + return s; +} + +// Parallel CSR run scan over one tile: row r is a run head iff +// dst[r] != dst[r - 1]; per-warp head ballots give every row its run index +// as a popcount prefix. All threads must call; barriers inside. +template +DEV_INLINE void scan_runs(int tid, + int rows, + const int* dst, + unsigned* head_masks, + int* run_node, + int* run_begin, + int* run_of, + int* n_runs) { + constexpr int TILE = NW * 32; // edges per tile + const int r = tid; + bool head = false; + if (r < TILE) { + head = r < rows && (r == 0 || dst[r] != dst[r - 1]); + } + if (tid < TILE) { + const unsigned m = __ballot_sync(0xffffffffu, head); + if ((tid & 31) == 0) { + head_masks[tid >> 5] = m; + } + } + __syncthreads(); + if (r < TILE) { + const int w = r >> 5, lane = r & 31; + int idx = 0; // heads strictly before row r +#pragma unroll + for (int i = 0; i < NW; ++i) { + if (i < w) { + idx += __popc(head_masks[i]); + } + } + idx += __popc(head_masks[w] & ((1u << lane) - 1u)); + if (r < rows) { + // A head row's exclusive count equals its run index; a non-head row's + // run head lies before it, so the exclusive count overshoots by one. + run_of[r] = head ? idx : idx - 1; + if (head) { + run_node[idx] = dst[r]; + run_begin[idx] = r; + } + } else { + run_of[r] = 0; + } + if (r == 0) { + int total = 0; +#pragma unroll + for (int i = 0; i < NW; ++i) { + total += __popc(head_masks[i]); + } + *n_runs = total; + run_begin[total] = rows; + } + } + __syncthreads(); +} + +// Per-edge tile tables shared by the forward and backward kernels; the tile +// width (edges per tile) is a template parameter so a forward can keep the +// 128-edge tile while its backward uses a narrower tile (fewer edges per +// thread, so the per-thread register footprint drops below the spill wall). +template +struct EdgeTablesT { + float rr[TILE][4]; // env-mat row premultiplied by mask / nnei + float radial[TILE]; // rr0 (unmasked MLP / table input) + float sw[TILE]; // raw switch value (strip gate factor) + int pair_idx[TILE]; + int dst[TILE]; + int run_node[TILE]; + int run_begin[TILE + 1]; + int run_of[TILE]; + unsigned head_masks[TILE / 32]; + int n_runs; +}; + +template +DEV_INLINE void stage_tile(int tid, + long tile_base, + int rows, + long n_edge, + int ntypes, + int one_side, + float rcut, + float rcut_smth, + float protection, + float inv_nnei, + const float* __restrict__ edge_vec, + const long* __restrict__ edge_index, + const bool* __restrict__ edge_mask, + const long* __restrict__ atype, + const float* __restrict__ davg, + const float* __restrict__ inv_dstd, + const int* __restrict__ order, + EdgeTablesT& T) { + if (tid < TILE) { + if (tid < rows) { + const int e = order[tile_base + tid]; + const auto s = + stage_edge(e, n_edge, ntypes, one_side, rcut, rcut_smth, protection, + edge_vec, edge_index, edge_mask, atype, davg, inv_dstd); + const float mm = (s.valid ? 1.f : 0.f) * inv_nnei; + T.radial[tid] = s.r0; + T.rr[tid][0] = s.r0 * mm; + T.rr[tid][1] = s.r1 * mm; + T.rr[tid][2] = s.r2 * mm; + T.rr[tid][3] = s.r3 * mm; + T.sw[tid] = s.sw; + T.pair_idx[tid] = s.pair_idx; + T.dst[tid] = s.dst; + } else { + // Tail rows of a partial tile: finite MLP input, zero moment weight, + // sentinel dst excluded from every run. + T.radial[tid] = 0.f; + T.rr[tid][0] = 0.f; + T.rr[tid][1] = 0.f; + T.rr[tid][2] = 0.f; + T.rr[tid][3] = 0.f; + T.sw[tid] = 0.f; + T.pair_idx[tid] = 0; + T.dst[tid] = -1; + } + } + __syncthreads(); + scan_runs(tid, rows, T.dst, T.head_masks, T.run_node, T.run_begin, + T.run_of, &T.n_runs); +} + +// Type-pair gate factor of one edge and channel (strip mode): +// gate_eff = gate_table[pair, c] (* sw when smooth_type_embedding). +// The gathered scalar rides L1/L2 (edges of one run share dst but not the +// neighbor type, so the rows repeat within a few hundred distinct pairs). +DEV_INLINE float gate_factor(const float* __restrict__ gate_table, + int ng, + int pair_idx, + int c, + float sw, + int smooth) { + const float gate = __ldg(gate_table + (long)pair_idx * ng + c); + return smooth ? gate * sw : gate; +} + +// ====================================================================== +// Gram contraction: one CTA per node. +// grrg[n, i * axis + j] = sum_k gr[n, k, i] * gr[n, k, j] +// rot_mat[n, i, :] = gr[n, 1:4, i] +// plus the appended center type embedding when concat_tebd is set. +// ====================================================================== +__global__ void gram_kernel(int n_node, + int ng, + int axis, + int tebd_dim, + int concat_tebd, + const float* __restrict__ gr, + const float* __restrict__ type_embedding, + const long* __restrict__ atype, + float* __restrict__ grrg, + float* __restrict__ rot_mat) { + const int n = blockIdx.x; + extern __shared__ float s_gr[]; // (4, ng) + for (int t = threadIdx.x; t < 4 * ng; t += blockDim.x) { + s_gr[t] = gr[(long)n * 4 * ng + t]; + } + __syncthreads(); + const int out_dim = ng * axis + (concat_tebd ? tebd_dim : 0); + float* out = grrg + (long)n * out_dim; + for (int i = threadIdx.x; i < ng; i += blockDim.x) { + const float g0 = s_gr[0 * ng + i], g1 = s_gr[1 * ng + i]; + const float g2 = s_gr[2 * ng + i], g3 = s_gr[3 * ng + i]; + for (int j = 0; j < axis; ++j) { + out[i * axis + j] = g0 * s_gr[0 * ng + j] + g1 * s_gr[1 * ng + j] + + g2 * s_gr[2 * ng + j] + g3 * s_gr[3 * ng + j]; + } + if (rot_mat != nullptr) { + rot_mat[((long)n * ng + i) * 3 + 0] = g1; + rot_mat[((long)n * ng + i) * 3 + 1] = g2; + rot_mat[((long)n * ng + i) * 3 + 2] = g3; + } + } + if (concat_tebd) { + const float* te = type_embedding + atype[n] * tebd_dim; + for (int t = threadIdx.x; t < tebd_dim; t += blockDim.x) { + out[ng * axis + t] = te[t]; + } + } +} + +// Gram backward: with D = d(grrg) reshaped to (ng, axis) rows and R = d(rot), +// dgr[n, k, c] = sum_j D[c, j] gr[k, j] +// + (c < axis) sum_i D[i, c] gr[k, i] +// + (k >= 1) R[c, k - 1]. +// The concat tebd tail of d(grrg) is a constant feature and carries no +// gradient; grrg_stride skips it. +__global__ void gram_backward_kernel(int n_node, + int ng, + int axis, + int grrg_stride, + const float* __restrict__ d_grrg, + const float* __restrict__ d_rot, + const float* __restrict__ gr, + float* __restrict__ dgr) { + const int n = blockIdx.x; + extern __shared__ float sm[]; // [0, 4*ng): gr; [4*ng, ...): d_grrg row + float* s_gr = sm; + float* s_dg = sm + 4 * ng; + for (int t = threadIdx.x; t < 4 * ng; t += blockDim.x) { + s_gr[t] = gr[(long)n * 4 * ng + t]; + } + for (int t = threadIdx.x; t < ng * axis; t += blockDim.x) { + s_dg[t] = d_grrg[(long)n * grrg_stride + t]; + } + __syncthreads(); + for (int c = threadIdx.x; c < ng; c += blockDim.x) { + float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f, acc3 = 0.f; + for (int j = 0; j < axis; ++j) { + const float d = s_dg[c * axis + j]; + acc0 = fmaf(d, s_gr[0 * ng + j], acc0); + acc1 = fmaf(d, s_gr[1 * ng + j], acc1); + acc2 = fmaf(d, s_gr[2 * ng + j], acc2); + acc3 = fmaf(d, s_gr[3 * ng + j], acc3); + } + if (c < axis) { + for (int i = 0; i < ng; ++i) { + const float d = s_dg[i * axis + c]; + acc0 = fmaf(d, s_gr[0 * ng + i], acc0); + acc1 = fmaf(d, s_gr[1 * ng + i], acc1); + acc2 = fmaf(d, s_gr[2 * ng + i], acc2); + acc3 = fmaf(d, s_gr[3 * ng + i], acc3); + } + } + if (d_rot) { + acc1 += d_rot[((long)n * ng + c) * 3 + 0]; + acc2 += d_rot[((long)n * ng + c) * 3 + 1]; + acc3 += d_rot[((long)n * ng + c) * 3 + 2]; + } + dgr[((long)n * 4 + 0) * ng + c] = acc0; + dgr[((long)n * 4 + 1) * ng + c] = acc1; + dgr[((long)n * 4 + 2) * ng + c] = acc2; + dgr[((long)n * 4 + 3) * ng + c] = acc3; + } +} + +// ====================================================================== +// Host-side launch utilities +// ====================================================================== +const float* optional_ptr(const torch::Tensor& t) { + return t.defined() && t.numel() ? t.data_ptr() : nullptr; +} + +int persistent_grid(long n_edge, int ctas_per_sm, int tile = kTileEdges) { + const int sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + return (int)std::min((long)sms * ctas_per_sm, ceil_div(n_edge, tile)); +} + +// CSR edge ordering shared by the descriptor forward entry points (the +// backwards receive the order as a saved tensor). +torch::Tensor build_edge_order(const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + long n_edge, + long n_node, + cudaStream_t stream) { + TORCH_CHECK( + n_edge <= std::numeric_limits::max(), + "dpa1_graph_descriptor: edge count exceeds the int32 order limit"); + auto i32 = + torch::TensorOptions().dtype(torch::kInt32).device(edge_index.device()); + auto counts = torch::zeros({n_node + 1}, i32); + auto unsorted = torch::zeros({1}, i32); + const long* dst = edge_index.data_ptr() + n_edge; + edge_order_count_kernel<<>>( + n_edge, dst, edge_mask.data_ptr(), counts.data_ptr(), + unsorted.data_ptr()); + auto offsets = torch::empty({n_node + 1}, i32); + { + size_t temp_bytes = 0; + cub::DeviceScan::ExclusiveSum(nullptr, temp_bytes, counts.data_ptr(), + offsets.data_ptr(), n_node + 1, stream); + auto temp = + torch::empty({(long)temp_bytes}, torch::TensorOptions() + .dtype(torch::kUInt8) + .device(edge_index.device())); + cub::DeviceScan::ExclusiveSum(temp.data_ptr(), temp_bytes, + counts.data_ptr(), + offsets.data_ptr(), n_node + 1, stream); + } + auto cursor = torch::zeros({n_node}, i32); + auto order = torch::empty({n_edge}, i32); + edge_order_scatter_kernel<<>>( + n_edge, dst, offsets.data_ptr(), unsorted.data_ptr(), + cursor.data_ptr(), order.data_ptr()); + DPA1_CHECK_LAUNCH("dpa1_graph edge order"); + return order; +} + +} // namespace diff --git a/source/op/pt/dpa1_graph_compress.cu b/source/op/pt/dpa1_graph_compress.cu new file mode 100644 index 0000000000..c06009b2a4 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress.cu @@ -0,0 +1,1661 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Geometrically compressed DPA1 graph descriptor over destination CSR edges. +// +// A warp owns one center node. Widths from 16 through 64 use two 16-lane +// sub-warps on alternating edges; each lane evaluates one or more spline +// channels. Wider tables retain one edge per warp to bound register pressure. +// The node moment and its Gram contraction remain in the same kernel. +// +// The backward recomputes the inexpensive spline value/derivative, contracts +// the descriptor gradient into the four environment channels, and writes each +// edge gradient exactly once. It is inference-oriented (one backward); the +// registered Python autograd bridge continues to expose the edge-vector +// gradient for the level-1 graph path. +// +// Every specialization has balanced (two CTA/SM launch bound) and occupancy +// (four CTA/SM launch bound) resource variants. The first uncaptured call times +// 128- and 256-thread launches on a bounded node sample and caches the selected +// variant per device and workload class. Device-family defaults remain valid +// when timing is disabled or CUDA Graph capture is active. + +#include +#include +#include + +#include +#include +#include + +#include "dpa1_graph_compress_tuning.h" + +namespace { + +using deepmd::dpa1_compress_tuning::device_properties; +using deepmd::dpa1_compress_tuning::KernelDirection; +using deepmd::dpa1_compress_tuning::LaunchConfig; +using deepmd::dpa1_compress_tuning::ResourcePolicy; +using deepmd::dpa1_compress_tuning::select_launch_config; +using deepmd::dpa1_compress_tuning::TuningKey; +using deepmd::dpa1_compress_tuning::type_count_class; +using deepmd::dpa1_compress_tuning::workload_degree_class; +using deepmd::dpa1_compress_tuning::workload_size_class; + +constexpr int kThreads = 256; +constexpr int kWarpSize = 32; + +#define COMPRESS_CHECK_LAUNCH(name) \ + do { \ + const cudaError_t error = cudaGetLastError(); \ + TORCH_CHECK(error == cudaSuccess, name, ": ", cudaGetErrorString(error)); \ + } while (0) + +__device__ __forceinline__ float switch_value(float radius, + float lower, + float upper) { + const float coordinate = + __fdividef(fminf(fmaxf(radius, lower), upper) - lower, upper - lower); + const float square = coordinate * coordinate; + return square * coordinate * (-6.0f * square + 15.0f * coordinate - 10.0f) + + 1.0f; +} + +__device__ __forceinline__ float switch_derivative(float radius, + float lower, + float upper) { + if (radius <= lower || radius >= upper) { + return 0.0f; + } + const float coordinate = __fdividef(radius - lower, upper - lower); + const float square = coordinate * coordinate; + return __fdividef( + -30.0f * square * square + 60.0f * square * coordinate - 30.0f * square, + upper - lower); +} + +struct TableLocation { + int index; + float coordinate; + float extrapolation; +}; + +__device__ __forceinline__ int high_tail_index( + float lower, float upper, float table_max, float stride0, float stride1) { + const float boundary = nextafterf(table_max, lower); + const int first_stride = static_cast(__fdividef(upper - lower, stride0)); + return first_stride + static_cast(__fdividef(boundary - upper, stride1)); +} + +__device__ __forceinline__ TableLocation locate_table(float radial, + float lower, + float upper, + float table_max, + float stride0, + float stride1) { + TableLocation location; + location.coordinate = radial; + location.extrapolation = 0.0f; + if (radial < lower) { + location.index = 0; + location.coordinate = 0.0f; + location.extrapolation = radial - lower; + } else if (radial < upper) { + location.index = static_cast(__fdividef(radial - lower, stride0)); + location.coordinate -= location.index * stride0 + lower; + } else if (radial < table_max) { + const int first_stride = + static_cast(__fdividef(upper - lower, stride0)); + location.index = + first_stride + static_cast(__fdividef(radial - upper, stride1)); + location.coordinate -= (location.index - first_stride) * stride1 + upper; + } else { + const int first_stride = + static_cast(__fdividef(upper - lower, stride0)); + location.index = high_tail_index(lower, upper, table_max, stride0, stride1); + location.coordinate = + table_max - ((location.index - first_stride) * stride1 + upper); + location.extrapolation = radial - table_max; + } + return location; +} + +__device__ __forceinline__ void load_coefficients(const float* table, + const TableLocation& location, + int channel, + int width, + float2& c01, + float2& c23, + float2& c45) { + const long offset = static_cast(location.index) * width * 6 + + static_cast(channel) * 6; + c01 = __ldg(reinterpret_cast(table + offset)); + c23 = __ldg(reinterpret_cast(table + offset + 2)); + c45 = __ldg(reinterpret_cast(table + offset + 4)); +} + +__device__ __forceinline__ float evaluate_table_forward( + const float* table, const TableLocation& location, int channel, int width) { + float2 c01, c23, c45; + load_coefficients(table, location, channel, width, c01, c23, c45); + const float value = + c01.x + (c01.y + (c23.x + (c23.y + (c45.x + c45.y * location.coordinate) * + location.coordinate) * + location.coordinate) * + location.coordinate) * + location.coordinate; + if (location.extrapolation == 0.0f) { + return value; + } + const float derivative = + c01.y + + (2.0f * c23.x + + (3.0f * c23.y + (4.0f * c45.x + 5.0f * c45.y * location.coordinate) * + location.coordinate) * + location.coordinate) * + location.coordinate; + return value + derivative * location.extrapolation; +} + +__device__ __forceinline__ float2 evaluate_table_backward( + const float* table, const TableLocation& location, int channel, int width) { + float2 c01, c23, c45; + load_coefficients(table, location, channel, width, c01, c23, c45); + float value = c45.y; + float derivative = 0.0f; + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c45.x); + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c23.y); + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c23.x); + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c01.y); + derivative = fmaf(derivative, location.coordinate, value); + value = fmaf(value, location.coordinate, c01.x); + return make_float2(value + derivative * location.extrapolation, derivative); +} + +struct EdgeEnvironment { + float radial; + float r0; + float r1; + float r2; + float r3; + float switch_factor; + float x; + float y; + float z; + float radius; + int pair_index; +}; + +template +__device__ __forceinline__ EdgeEnvironment +load_environment(long edge, + int center_type, + int ntypes, + bool one_side, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + const float* edge_vec, + const index_t* edge_index, + const long* atype, + const float* average, + const float* inverse_stddev) { + EdgeEnvironment environment; + const long source = static_cast(edge_index[edge]); + const int neighbor_type = static_cast(atype[source]); + environment.x = edge_vec[edge * 3 + 0]; + environment.y = edge_vec[edge * 3 + 1]; + environment.z = edge_vec[edge * 3 + 2]; + const float square_length = environment.x * environment.x + + environment.y * environment.y + + environment.z * environment.z; + environment.radius = + square_length > 0.0f ? square_length * rsqrtf(square_length) : 0.0f; + const float denominator = environment.radius + protection; + environment.switch_factor = + switch_value(environment.radius, rcut_smooth, rcut); + const float inverse_radius = __fdividef(1.0f, denominator); + const float radial_scale = + environment.switch_factor * inverse_radius * inverse_radius; + const float* center_average = average + static_cast(center_type) * 4; + const float* center_inverse_stddev = + inverse_stddev + static_cast(center_type) * 4; + environment.radial = + (environment.switch_factor * inverse_radius - center_average[0]) * + center_inverse_stddev[0]; + environment.r0 = environment.radial * inverse_neighbors; + environment.r1 = (environment.x * radial_scale - center_average[1]) * + center_inverse_stddev[1] * inverse_neighbors; + environment.r2 = (environment.y * radial_scale - center_average[2]) * + center_inverse_stddev[2] * inverse_neighbors; + environment.r3 = (environment.z * radial_scale - center_average[3]) * + center_inverse_stddev[3] * inverse_neighbors; + environment.pair_index = + one_side ? neighbor_type : center_type * ntypes + neighbor_type; + return environment; +} + +__device__ __forceinline__ EdgeEnvironment +broadcast_environment(EdgeEnvironment value, int source_lane, unsigned mask) { + value.radial = __shfl_sync(mask, value.radial, source_lane); + value.r0 = __shfl_sync(mask, value.r0, source_lane); + value.r1 = __shfl_sync(mask, value.r1, source_lane); + value.r2 = __shfl_sync(mask, value.r2, source_lane); + value.r3 = __shfl_sync(mask, value.r3, source_lane); + value.switch_factor = __shfl_sync(mask, value.switch_factor, source_lane); + value.x = __shfl_sync(mask, value.x, source_lane); + value.y = __shfl_sync(mask, value.y, source_lane); + value.z = __shfl_sync(mask, value.z, source_lane); + value.radius = __shfl_sync(mask, value.radius, source_lane); + value.pair_index = __shfl_sync(mask, value.pair_index, source_lane); + return value; +} + +__device__ __forceinline__ TableLocation broadcast_location(TableLocation value, + int source_lane, + unsigned mask) { + value.index = __shfl_sync(mask, value.index, source_lane); + value.coordinate = __shfl_sync(mask, value.coordinate, source_lane); + value.extrapolation = __shfl_sync(mask, value.extrapolation, source_lane); + return value; +} + +__device__ __forceinline__ void store_edge_gradient( + long edge, + const EdgeEnvironment& environment, + float partial0, + float partial1, + float partial2, + float partial3, + float partial_radial, + float partial_switch, + float inverse_neighbors, + float inverse_stddev0, + float inverse_stddev1, + float inverse_stddev2, + float inverse_stddev3, + float rcut, + float rcut_smooth, + float protection, + float* edge_gradient) { + const float inverse_denominator = + __fdividef(1.0f, environment.radius + protection); + const float inverse_length = + environment.radius > 0.0f ? __fdividef(1.0f, environment.radius) : 0.0f; + const float switch_gradient = + switch_derivative(environment.radius, rcut_smooth, rcut); + const float gradient_radial = + (partial0 * inverse_neighbors + partial_radial) * inverse_stddev0; + const float gradient_x = partial1 * inverse_neighbors * inverse_stddev1; + const float gradient_y = partial2 * inverse_neighbors * inverse_stddev2; + const float gradient_z = partial3 * inverse_neighbors * inverse_stddev3; + const float directional = gradient_x * environment.x + + gradient_y * environment.y + + gradient_z * environment.z; + const float coefficient = + (gradient_radial * inverse_denominator * + (switch_gradient - environment.switch_factor * inverse_denominator) + + directional * inverse_denominator * inverse_denominator * + (switch_gradient - + 2.0f * environment.switch_factor * inverse_denominator) + + partial_switch * switch_gradient) * + inverse_length; + const float vector_scale = + environment.switch_factor * inverse_denominator * inverse_denominator; + edge_gradient[edge * 3 + 0] = + coefficient * environment.x + vector_scale * gradient_x; + edge_gradient[edge * 3 + 1] = + coefficient * environment.y + vector_scale * gradient_y; + edge_gradient[edge * 3 + 2] = + coefficient * environment.z + vector_scale * gradient_z; +} + +template +struct ChannelPolicy { + static constexpr bool use_half_warp = Width >= 16 && Width <= 64; + static constexpr int accumulation_groups = + use_half_warp ? Width / 16 : (Width + 31) / 32; + static constexpr int gradient_groups = (Width + 31) / 32; +}; + +template +__device__ __forceinline__ long edge_at_csr_position( + long position, const index_t* destination_order) { + if constexpr (Canonical) { + return position; + } + return static_cast(destination_order[position]); +} + +template +__device__ __forceinline__ bool edge_is_active(long edge, + const bool* edge_mask) { + if constexpr (Masked) { + return edge_mask[edge]; + } + return true; +} + +template +__global__ +__launch_bounds__(kThreads, MinimumBlocks) void compressed_forward_kernel( + long node_count, + int ntypes, + bool one_side, + bool smooth, + int axis, + bool concatenate_type_embedding, + bool write_rotation, + int type_embedding_dim, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const float* __restrict__ edge_vec, + const index_t* __restrict__ edge_index, + const bool* __restrict__ edge_mask, + const index_t* __restrict__ destination_order, + const long* __restrict__ destination_row_ptr, + const long* __restrict__ atype, + const float* __restrict__ type_embedding, + const float* __restrict__ average, + const float* __restrict__ inverse_stddev, + const float* __restrict__ table, + const float* __restrict__ gate_table, + float* __restrict__ descriptor, + float* __restrict__ rotation, + float* __restrict__ moment) { + constexpr unsigned kWarpMask = 0xffffffffu; + constexpr int kGroups = ChannelPolicy::accumulation_groups; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int warps_per_block = blockDim.x / kWarpSize; + const long node = static_cast(blockIdx.x) * warps_per_block + warp; + if (node >= node_count) { + return; + } + + float accumulator0[kGroups] = {}; + float accumulator1[kGroups] = {}; + float accumulator2[kGroups] = {}; + float accumulator3[kGroups] = {}; + int center_type = lane == 0 ? static_cast(atype[node]) : 0; + center_type = __shfl_sync(kWarpMask, center_type, 0); + const long begin = destination_row_ptr[node]; + const long end = destination_row_ptr[node + 1]; + + if constexpr (ChannelPolicy::use_half_warp) { + const int half = lane >> 4; + const int half_lane = lane & 15; + const int leader = half * 16; + const unsigned half_mask = half == 0 ? 0x0000ffffu : 0xffff0000u; + for (long position = begin + half; position < end; position += 2) { + const long edge = + edge_at_csr_position(position, destination_order); + if (!edge_is_active(edge, edge_mask)) { + continue; + } + EdgeEnvironment environment{}; + TableLocation location{}; + if (half_lane == 0) { + environment = load_environment(edge, center_type, ntypes, one_side, + rcut, rcut_smooth, protection, + inverse_neighbors, edge_vec, edge_index, + atype, average, inverse_stddev); + location = locate_table(environment.radial, lower, upper, table_max, + stride0, stride1); + } + environment = broadcast_environment(environment, leader, half_mask); + location = broadcast_location(location, leader, half_mask); +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 16 + half_lane; + const float table_value = + evaluate_table_forward(table, location, channel, Width); + const float gate = + __ldg(gate_table + + static_cast(environment.pair_index) * Width + channel); + const float effective_gate = + smooth ? gate * environment.switch_factor : gate; + const float embedding = table_value * (1.0f + effective_gate); + accumulator0[group] = + fmaf(environment.r0, embedding, accumulator0[group]); + accumulator1[group] = + fmaf(environment.r1, embedding, accumulator1[group]); + accumulator2[group] = + fmaf(environment.r2, embedding, accumulator2[group]); + accumulator3[group] = + fmaf(environment.r3, embedding, accumulator3[group]); + } + } +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + accumulator0[group] += + __shfl_xor_sync(kWarpMask, accumulator0[group], 16); + accumulator1[group] += + __shfl_xor_sync(kWarpMask, accumulator1[group], 16); + accumulator2[group] += + __shfl_xor_sync(kWarpMask, accumulator2[group], 16); + accumulator3[group] += + __shfl_xor_sync(kWarpMask, accumulator3[group], 16); + } + } else { + for (long position = begin; position < end; ++position) { + const long edge = + edge_at_csr_position(position, destination_order); + if (!edge_is_active(edge, edge_mask)) { + continue; + } + EdgeEnvironment environment{}; + TableLocation location{}; + if (lane == 0) { + environment = load_environment(edge, center_type, ntypes, one_side, + rcut, rcut_smooth, protection, + inverse_neighbors, edge_vec, edge_index, + atype, average, inverse_stddev); + location = locate_table(environment.radial, lower, upper, table_max, + stride0, stride1); + } + environment = broadcast_environment(environment, 0, kWarpMask); + location = broadcast_location(location, 0, kWarpMask); +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 32 + lane; + if (channel < Width) { + const float table_value = + evaluate_table_forward(table, location, channel, Width); + const float gate = __ldg( + gate_table + static_cast(environment.pair_index) * Width + + channel); + const float effective_gate = + smooth ? gate * environment.switch_factor : gate; + const float embedding = table_value * (1.0f + effective_gate); + accumulator0[group] = + fmaf(environment.r0, embedding, accumulator0[group]); + accumulator1[group] = + fmaf(environment.r1, embedding, accumulator1[group]); + accumulator2[group] = + fmaf(environment.r2, embedding, accumulator2[group]); + accumulator3[group] = + fmaf(environment.r3, embedding, accumulator3[group]); + } + } + } + } + + const long moment_base = node * 4 * Width; + if constexpr (ChannelPolicy::use_half_warp) { + if (lane < 16) { +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 16 + lane; + moment[moment_base + 0 * Width + channel] = accumulator0[group]; + moment[moment_base + 1 * Width + channel] = accumulator1[group]; + moment[moment_base + 2 * Width + channel] = accumulator2[group]; + moment[moment_base + 3 * Width + channel] = accumulator3[group]; + } + } + } else { +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 32 + lane; + if (channel < Width) { + moment[moment_base + 0 * Width + channel] = accumulator0[group]; + moment[moment_base + 1 * Width + channel] = accumulator1[group]; + moment[moment_base + 2 * Width + channel] = accumulator2[group]; + moment[moment_base + 3 * Width + channel] = accumulator3[group]; + } + } + } + + const int output_dim = + Width * axis + (concatenate_type_embedding ? type_embedding_dim : 0); + float* output = descriptor + node * output_dim; + if constexpr (ChannelPolicy::use_half_warp) { +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 16 + (lane & 15); + for (int axis_channel = 0; axis_channel < axis; ++axis_channel) { + const float axis0 = + __shfl_sync(kWarpMask, accumulator0[0], axis_channel); + const float axis1 = + __shfl_sync(kWarpMask, accumulator1[0], axis_channel); + const float axis2 = + __shfl_sync(kWarpMask, accumulator2[0], axis_channel); + const float axis3 = + __shfl_sync(kWarpMask, accumulator3[0], axis_channel); + if (lane < 16) { + output[channel * axis + axis_channel] = + accumulator0[group] * axis0 + accumulator1[group] * axis1 + + accumulator2[group] * axis2 + accumulator3[group] * axis3; + } + } + if (write_rotation && lane < 16) { + float* rotation_row = rotation + (node * Width + channel) * 3; + rotation_row[0] = accumulator1[group]; + rotation_row[1] = accumulator2[group]; + rotation_row[2] = accumulator3[group]; + } + } + } else { +#pragma unroll + for (int group = 0; group < kGroups; ++group) { + const int channel = group * 32 + lane; + for (int axis_channel = 0; axis_channel < axis; ++axis_channel) { + const float axis0 = + __shfl_sync(kWarpMask, accumulator0[0], axis_channel); + const float axis1 = + __shfl_sync(kWarpMask, accumulator1[0], axis_channel); + const float axis2 = + __shfl_sync(kWarpMask, accumulator2[0], axis_channel); + const float axis3 = + __shfl_sync(kWarpMask, accumulator3[0], axis_channel); + if (channel < Width) { + output[channel * axis + axis_channel] = + accumulator0[group] * axis0 + accumulator1[group] * axis1 + + accumulator2[group] * axis2 + accumulator3[group] * axis3; + } + } + if (write_rotation && channel < Width) { + float* rotation_row = rotation + (node * Width + channel) * 3; + rotation_row[0] = accumulator1[group]; + rotation_row[1] = accumulator2[group]; + rotation_row[2] = accumulator3[group]; + } + } + } + if (concatenate_type_embedding) { + for (int channel = lane; channel < type_embedding_dim; channel += 32) { + output[Width * axis + channel] = + type_embedding[static_cast(center_type) * type_embedding_dim + + channel]; + } + } +} + +template +__global__ +__launch_bounds__(kThreads, MinimumBlocks) void compressed_backward_kernel( + long node_count, + long edge_count, + int ntypes, + bool one_side, + bool smooth, + int axis, + int descriptor_stride, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const float* __restrict__ descriptor_gradient, + const float* __restrict__ rotation_gradient, + const float* __restrict__ moment, + const float* __restrict__ edge_vec, + const index_t* __restrict__ edge_index, + const bool* __restrict__ edge_mask, + const index_t* __restrict__ destination_order, + const long* __restrict__ destination_row_ptr, + const long* __restrict__ atype, + const float* __restrict__ average, + const float* __restrict__ inverse_stddev, + const float* __restrict__ table, + const float* __restrict__ gate_table, + float* __restrict__ edge_gradient) { + constexpr unsigned kWarpMask = 0xffffffffu; + constexpr int kGradientGroups = ChannelPolicy::gradient_groups; + constexpr int kEdgeGroups = ChannelPolicy::accumulation_groups; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int warps_per_block = blockDim.x / kWarpSize; + const long node = static_cast(blockIdx.x) * warps_per_block + warp; + if (node >= node_count) { + return; + } + + float gradient0[kGradientGroups] = {}; + float gradient1[kGradientGroups] = {}; + float gradient2[kGradientGroups] = {}; + float gradient3[kGradientGroups] = {}; + float own_moment0[kGradientGroups] = {}; + float own_moment1[kGradientGroups] = {}; + float own_moment2[kGradientGroups] = {}; + float own_moment3[kGradientGroups] = {}; + const long moment_base = node * 4 * Width; + const float* node_descriptor_gradient = + descriptor_gradient + node * descriptor_stride; + +#pragma unroll + for (int group = 0; group < kGradientGroups; ++group) { + const int channel = group * 32 + lane; + if (channel < Width) { + own_moment0[group] = __ldg(moment + moment_base + 0 * Width + channel); + own_moment1[group] = __ldg(moment + moment_base + 1 * Width + channel); + own_moment2[group] = __ldg(moment + moment_base + 2 * Width + channel); + own_moment3[group] = __ldg(moment + moment_base + 3 * Width + channel); + } + for (int axis_channel = 0; axis_channel < axis; ++axis_channel) { + const float axis0 = __shfl_sync(kWarpMask, own_moment0[0], axis_channel); + const float axis1 = __shfl_sync(kWarpMask, own_moment1[0], axis_channel); + const float axis2 = __shfl_sync(kWarpMask, own_moment2[0], axis_channel); + const float axis3 = __shfl_sync(kWarpMask, own_moment3[0], axis_channel); + if (channel < Width) { + const float value = + __ldg(node_descriptor_gradient + channel * axis + axis_channel); + gradient0[group] = fmaf(value, axis0, gradient0[group]); + gradient1[group] = fmaf(value, axis1, gradient1[group]); + gradient2[group] = fmaf(value, axis2, gradient2[group]); + gradient3[group] = fmaf(value, axis3, gradient3[group]); + } + } + if (channel < axis) { + for (int input = 0; input < Width; ++input) { + const float value = + __ldg(node_descriptor_gradient + input * axis + channel); + gradient0[group] = + fmaf(value, __ldg(moment + moment_base + 0 * Width + input), + gradient0[group]); + gradient1[group] = + fmaf(value, __ldg(moment + moment_base + 1 * Width + input), + gradient1[group]); + gradient2[group] = + fmaf(value, __ldg(moment + moment_base + 2 * Width + input), + gradient2[group]); + gradient3[group] = + fmaf(value, __ldg(moment + moment_base + 3 * Width + input), + gradient3[group]); + } + } + if (rotation_gradient != nullptr && channel < Width) { + const long rotation_offset = (node * Width + channel) * 3; + gradient1[group] += __ldg(rotation_gradient + rotation_offset + 0); + gradient2[group] += __ldg(rotation_gradient + rotation_offset + 1); + gradient3[group] += __ldg(rotation_gradient + rotation_offset + 2); + } + } + + int center_type = lane == 0 ? static_cast(atype[node]) : 0; + center_type = __shfl_sync(kWarpMask, center_type, 0); + const float inverse_stddev0 = + __ldg(inverse_stddev + static_cast(center_type) * 4 + 0); + const float inverse_stddev1 = + __ldg(inverse_stddev + static_cast(center_type) * 4 + 1); + const float inverse_stddev2 = + __ldg(inverse_stddev + static_cast(center_type) * 4 + 2); + const float inverse_stddev3 = + __ldg(inverse_stddev + static_cast(center_type) * 4 + 3); + const long begin = destination_row_ptr[node]; + const long end = destination_row_ptr[node + 1]; + + if constexpr (ChannelPolicy::use_half_warp) { + const int half = lane >> 4; + const int half_lane = lane & 15; + const int leader = half * 16; + const unsigned half_mask = half == 0 ? 0x0000ffffu : 0xffff0000u; + float edge_gradient0[kEdgeGroups] = {}; + float edge_gradient1[kEdgeGroups] = {}; + float edge_gradient2[kEdgeGroups] = {}; + float edge_gradient3[kEdgeGroups] = {}; +#pragma unroll + for (int group = 0; group < kEdgeGroups; ++group) { + const int channel = group * 16 + half_lane; + const int owner_group = channel / 32; + const int owner_lane = channel & 31; + edge_gradient0[group] = + __shfl_sync(kWarpMask, gradient0[owner_group], owner_lane); + edge_gradient1[group] = + __shfl_sync(kWarpMask, gradient1[owner_group], owner_lane); + edge_gradient2[group] = + __shfl_sync(kWarpMask, gradient2[owner_group], owner_lane); + edge_gradient3[group] = + __shfl_sync(kWarpMask, gradient3[owner_group], owner_lane); + } + + for (long position = begin + half; position < end; position += 2) { + const long edge = + edge_at_csr_position(position, destination_order); + if (!edge_is_active(edge, edge_mask)) { + if (half_lane == 0) { + edge_gradient[edge * 3 + 0] = 0.0f; + edge_gradient[edge * 3 + 1] = 0.0f; + edge_gradient[edge * 3 + 2] = 0.0f; + } + continue; + } + EdgeEnvironment environment{}; + TableLocation location{}; + if (half_lane == 0) { + environment = load_environment(edge, center_type, ntypes, one_side, + rcut, rcut_smooth, protection, + inverse_neighbors, edge_vec, edge_index, + atype, average, inverse_stddev); + location = locate_table(environment.radial, lower, upper, table_max, + stride0, stride1); + } + environment = broadcast_environment(environment, leader, half_mask); + location = broadcast_location(location, leader, half_mask); + float partial0 = 0.0f; + float partial1 = 0.0f; + float partial2 = 0.0f; + float partial3 = 0.0f; + float partial_radial = 0.0f; + float partial_switch = 0.0f; +#pragma unroll + for (int group = 0; group < kEdgeGroups; ++group) { + const int channel = group * 16 + half_lane; + const float d0 = edge_gradient0[group]; + const float d1 = edge_gradient1[group]; + const float d2 = edge_gradient2[group]; + const float d3 = edge_gradient3[group]; + const float descriptor_product = + environment.r0 * d0 + environment.r1 * d1 + environment.r2 * d2 + + environment.r3 * d3; + const float2 table_value = + evaluate_table_backward(table, location, channel, Width); + const float gate = + __ldg(gate_table + + static_cast(environment.pair_index) * Width + channel); + const float effective_gate = + smooth ? gate * environment.switch_factor : gate; + const float embedding = table_value.x * (1.0f + effective_gate); + if (smooth) { + partial_switch = + fmaf(descriptor_product * table_value.x, gate, partial_switch); + } + partial_radial = fmaf(descriptor_product * (1.0f + effective_gate), + table_value.y, partial_radial); + partial0 = fmaf(embedding, d0, partial0); + partial1 = fmaf(embedding, d1, partial1); + partial2 = fmaf(embedding, d2, partial2); + partial3 = fmaf(embedding, d3, partial3); + } +#pragma unroll + for (int offset = 8; offset > 0; offset >>= 1) { + partial0 += __shfl_down_sync(half_mask, partial0, offset, 16); + partial1 += __shfl_down_sync(half_mask, partial1, offset, 16); + partial2 += __shfl_down_sync(half_mask, partial2, offset, 16); + partial3 += __shfl_down_sync(half_mask, partial3, offset, 16); + partial_radial += + __shfl_down_sync(half_mask, partial_radial, offset, 16); + partial_switch += + __shfl_down_sync(half_mask, partial_switch, offset, 16); + } + if (half_lane == 0) { + store_edge_gradient(edge, environment, partial0, partial1, partial2, + partial3, partial_radial, partial_switch, + inverse_neighbors, inverse_stddev0, inverse_stddev1, + inverse_stddev2, inverse_stddev3, rcut, rcut_smooth, + protection, edge_gradient); + } + } + } else { + for (long position = begin; position < end; ++position) { + const long edge = + edge_at_csr_position(position, destination_order); + if (!edge_is_active(edge, edge_mask)) { + if (lane == 0) { + edge_gradient[edge * 3 + 0] = 0.0f; + edge_gradient[edge * 3 + 1] = 0.0f; + edge_gradient[edge * 3 + 2] = 0.0f; + } + continue; + } + EdgeEnvironment environment{}; + TableLocation location{}; + if (lane == 0) { + environment = load_environment(edge, center_type, ntypes, one_side, + rcut, rcut_smooth, protection, + inverse_neighbors, edge_vec, edge_index, + atype, average, inverse_stddev); + location = locate_table(environment.radial, lower, upper, table_max, + stride0, stride1); + } + environment = broadcast_environment(environment, 0, kWarpMask); + location = broadcast_location(location, 0, kWarpMask); + float partial0 = 0.0f; + float partial1 = 0.0f; + float partial2 = 0.0f; + float partial3 = 0.0f; + float partial_radial = 0.0f; + float partial_switch = 0.0f; +#pragma unroll + for (int group = 0; group < kGradientGroups; ++group) { + const int channel = group * 32 + lane; + if (channel < Width) { + const float d0 = gradient0[group]; + const float d1 = gradient1[group]; + const float d2 = gradient2[group]; + const float d3 = gradient3[group]; + const float descriptor_product = + environment.r0 * d0 + environment.r1 * d1 + environment.r2 * d2 + + environment.r3 * d3; + const float2 table_value = + evaluate_table_backward(table, location, channel, Width); + const float gate = __ldg( + gate_table + static_cast(environment.pair_index) * Width + + channel); + const float effective_gate = + smooth ? gate * environment.switch_factor : gate; + const float embedding = table_value.x * (1.0f + effective_gate); + if (smooth) { + partial_switch = + fmaf(descriptor_product * table_value.x, gate, partial_switch); + } + partial_radial = fmaf(descriptor_product * (1.0f + effective_gate), + table_value.y, partial_radial); + partial0 = fmaf(embedding, d0, partial0); + partial1 = fmaf(embedding, d1, partial1); + partial2 = fmaf(embedding, d2, partial2); + partial3 = fmaf(embedding, d3, partial3); + } + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + partial0 += __shfl_down_sync(kWarpMask, partial0, offset); + partial1 += __shfl_down_sync(kWarpMask, partial1, offset); + partial2 += __shfl_down_sync(kWarpMask, partial2, offset); + partial3 += __shfl_down_sync(kWarpMask, partial3, offset); + partial_radial += __shfl_down_sync(kWarpMask, partial_radial, offset); + partial_switch += __shfl_down_sync(kWarpMask, partial_switch, offset); + } + if (lane == 0) { + store_edge_gradient(edge, environment, partial0, partial1, partial2, + partial3, partial_radial, partial_switch, + inverse_neighbors, inverse_stddev0, inverse_stddev1, + inverse_stddev2, inverse_stddev3, rcut, rcut_smooth, + protection, edge_gradient); + } + } + } +} + +template +__global__ void zero_padding_kernel( + long node_count, + long edge_count, + const index_t* __restrict__ destination_order, + const long* __restrict__ destination_row_ptr, + float* __restrict__ edge_gradient) { + const long valid_edge_count = destination_row_ptr[node_count]; + for (long position = valid_edge_count + blockIdx.x * blockDim.x + threadIdx.x; + position < edge_count; + position += static_cast(blockDim.x) * gridDim.x) { + const long edge = + edge_at_csr_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; + } +} + +template +void launch_forward_variant(long node_count, + int threads, + int ntypes, + bool one_side, + bool smooth, + int axis, + bool concatenate_type_embedding, + bool write_rotation, + int type_embedding_dim, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& atype, + const torch::Tensor& type_embedding, + const torch::Tensor& average, + const torch::Tensor& inverse_stddev, + const torch::Tensor& table, + const torch::Tensor& gate_table, + torch::Tensor& descriptor, + torch::Tensor& rotation, + torch::Tensor& moment, + cudaStream_t stream) { + const int warps_per_block = threads / kWarpSize; + const int blocks = + static_cast((node_count + warps_per_block - 1) / warps_per_block); + compressed_forward_kernel + <<>>( + node_count, ntypes, one_side, smooth, axis, + concatenate_type_embedding, write_rotation, type_embedding_dim, rcut, + rcut_smooth, protection, inverse_neighbors, lower, upper, table_max, + stride0, stride1, edge_vec.data_ptr(), + edge_index.data_ptr(), + edge_mask.numel() ? edge_mask.data_ptr() : nullptr, + destination_order.numel() ? destination_order.data_ptr() + : nullptr, + destination_row_ptr.data_ptr(), atype.data_ptr(), + type_embedding.data_ptr(), average.data_ptr(), + inverse_stddev.data_ptr(), table.data_ptr(), + gate_table.data_ptr(), descriptor.data_ptr(), + write_rotation ? rotation.data_ptr() : nullptr, + moment.data_ptr()); +} + +template +void launch_forward(long node_count, + int ntypes, + bool one_side, + bool smooth, + int axis, + bool concatenate_type_embedding, + bool write_rotation, + int type_embedding_dim, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& atype, + const torch::Tensor& type_embedding, + const torch::Tensor& average, + const torch::Tensor& inverse_stddev, + const torch::Tensor& table, + const torch::Tensor& gate_table, + torch::Tensor& descriptor, + torch::Tensor& rotation, + torch::Tensor& moment, + cudaStream_t stream) { + const int device = edge_vec.get_device(); + const cudaDeviceProp& properties = device_properties(device); + const TuningKey key = { + device, + static_cast(KernelDirection::kForward), + Width, + axis, + Canonical ? 1 : 0, + static_cast(sizeof(index_t)), + (one_side ? 1 : 0) | (smooth ? 2 : 0) | + (concatenate_type_embedding ? 4 : 0) | (write_rotation ? 8 : 0) | + (Masked ? 16 : 0), + concatenate_type_embedding ? type_embedding_dim : 0, + type_count_class(ntypes), + workload_size_class(node_count, properties.multiProcessorCount), + workload_degree_class(node_count, edge_vec.size(0)), + }; + const auto launch = [&](const LaunchConfig& config, long count) { + if (config.resource == ResourcePolicy::kOccupancy) { + launch_forward_variant( + count, config.threads, ntypes, one_side, smooth, axis, + concatenate_type_embedding, write_rotation, type_embedding_dim, rcut, + rcut_smooth, protection, inverse_neighbors, lower, upper, table_max, + stride0, stride1, edge_vec, edge_index, edge_mask, destination_order, + destination_row_ptr, atype, type_embedding, average, inverse_stddev, + table, gate_table, descriptor, rotation, moment, stream); + } else { + launch_forward_variant( + count, config.threads, ntypes, one_side, smooth, axis, + concatenate_type_embedding, write_rotation, type_embedding_dim, rcut, + rcut_smooth, protection, inverse_neighbors, lower, upper, table_max, + stride0, stride1, edge_vec, edge_index, edge_mask, destination_order, + destination_row_ptr, atype, type_embedding, average, inverse_stddev, + table, gate_table, descriptor, rotation, moment, stream); + } + }; + const LaunchConfig config = + select_launch_config(key, properties, node_count, stream, launch); + launch(config, node_count); + COMPRESS_CHECK_LAUNCH("dpa1_graph_compress forward"); +} + +template +void launch_backward_variant(long node_count, + int threads, + long edge_count, + int ntypes, + bool one_side, + bool smooth, + int axis, + int descriptor_stride, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const torch::Tensor& descriptor_gradient, + const float* rotation_gradient, + const torch::Tensor& moment, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& atype, + const torch::Tensor& average, + const torch::Tensor& inverse_stddev, + const torch::Tensor& table, + const torch::Tensor& gate_table, + torch::Tensor& edge_gradient, + cudaStream_t stream) { + const int warps_per_block = threads / kWarpSize; + const int blocks = + static_cast((node_count + warps_per_block - 1) / warps_per_block); + compressed_backward_kernel + <<>>( + node_count, edge_count, ntypes, one_side, smooth, axis, + descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, + lower, upper, table_max, stride0, stride1, + descriptor_gradient.data_ptr(), rotation_gradient, + moment.data_ptr(), edge_vec.data_ptr(), + edge_index.data_ptr(), + edge_mask.numel() ? edge_mask.data_ptr() : nullptr, + destination_order.numel() ? destination_order.data_ptr() + : nullptr, + destination_row_ptr.data_ptr(), atype.data_ptr(), + average.data_ptr(), inverse_stddev.data_ptr(), + table.data_ptr(), gate_table.data_ptr(), + edge_gradient.data_ptr()); +} + +template +void launch_backward(long node_count, + long edge_count, + int ntypes, + bool one_side, + bool smooth, + int axis, + int descriptor_stride, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const torch::Tensor& descriptor_gradient, + const float* rotation_gradient, + const torch::Tensor& moment, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& atype, + const torch::Tensor& average, + const torch::Tensor& inverse_stddev, + const torch::Tensor& table, + const torch::Tensor& gate_table, + torch::Tensor& edge_gradient, + cudaStream_t stream) { + const int device = edge_vec.get_device(); + const cudaDeviceProp& properties = device_properties(device); + const TuningKey key = { + device, + static_cast(KernelDirection::kBackward), + Width, + axis, + Canonical ? 1 : 0, + static_cast(sizeof(index_t)), + (one_side ? 1 : 0) | (smooth ? 2 : 0) | + (rotation_gradient != nullptr ? 8 : 0) | (Masked ? 16 : 0), + descriptor_stride, + type_count_class(ntypes), + workload_size_class(node_count, properties.multiProcessorCount), + workload_degree_class(node_count, edge_count), + }; + const auto launch = [&](const LaunchConfig& config, long count) { + if (config.resource == ResourcePolicy::kOccupancy) { + launch_backward_variant( + count, config.threads, edge_count, ntypes, one_side, smooth, axis, + descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, + lower, upper, table_max, stride0, stride1, descriptor_gradient, + rotation_gradient, moment, edge_vec, edge_index, edge_mask, + destination_order, destination_row_ptr, atype, average, + inverse_stddev, table, gate_table, edge_gradient, stream); + } else { + launch_backward_variant( + count, config.threads, edge_count, ntypes, one_side, smooth, axis, + descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, + lower, upper, table_max, stride0, stride1, descriptor_gradient, + rotation_gradient, moment, edge_vec, edge_index, edge_mask, + destination_order, destination_row_ptr, atype, average, + inverse_stddev, table, gate_table, edge_gradient, stream); + } + }; + const LaunchConfig config = + select_launch_config(key, properties, node_count, stream, launch); + launch(config, node_count); + COMPRESS_CHECK_LAUNCH("dpa1_graph_compress backward"); + zero_padding_kernel<<<1, kThreads, 0, stream>>>( + node_count, edge_count, + destination_order.numel() ? destination_order.data_ptr() + : nullptr, + destination_row_ptr.data_ptr(), edge_gradient.data_ptr()); + COMPRESS_CHECK_LAUNCH("dpa1_graph_compress padding"); +} + +template +void dispatch_forward(int width, + long node_count, + int ntypes, + bool one_side, + bool smooth, + int axis, + bool canonical, + bool masked, + bool concatenate_type_embedding, + bool write_rotation, + int type_embedding_dim, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& atype, + const torch::Tensor& type_embedding, + const torch::Tensor& average, + const torch::Tensor& inverse_stddev, + const torch::Tensor& table, + const torch::Tensor& gate_table, + torch::Tensor& descriptor, + torch::Tensor& rotation, + torch::Tensor& moment, + cudaStream_t stream) { +#define DISPATCH_WIDTH(value) \ + if (width == value) { \ + if (canonical && !masked) { \ + launch_forward( \ + node_count, ntypes, one_side, smooth, axis, \ + concatenate_type_embedding, write_rotation, type_embedding_dim, \ + rcut, rcut_smooth, protection, inverse_neighbors, lower, upper, \ + table_max, stride0, stride1, edge_vec, edge_index, edge_mask, \ + destination_order, destination_row_ptr, atype, type_embedding, \ + average, inverse_stddev, table, gate_table, descriptor, rotation, \ + moment, stream); \ + } else if (canonical) { \ + launch_forward( \ + node_count, ntypes, one_side, smooth, axis, \ + concatenate_type_embedding, write_rotation, type_embedding_dim, \ + rcut, rcut_smooth, protection, inverse_neighbors, lower, upper, \ + table_max, stride0, stride1, edge_vec, edge_index, edge_mask, \ + destination_order, destination_row_ptr, atype, type_embedding, \ + average, inverse_stddev, table, gate_table, descriptor, rotation, \ + moment, stream); \ + } else { \ + launch_forward( \ + node_count, ntypes, one_side, smooth, axis, \ + concatenate_type_embedding, write_rotation, type_embedding_dim, \ + rcut, rcut_smooth, protection, inverse_neighbors, lower, upper, \ + table_max, stride0, stride1, edge_vec, edge_index, edge_mask, \ + destination_order, destination_row_ptr, atype, type_embedding, \ + average, inverse_stddev, table, gate_table, descriptor, rotation, \ + moment, stream); \ + } \ + return; \ + } + DISPATCH_WIDTH(8) + DISPATCH_WIDTH(16) + DISPATCH_WIDTH(32) + DISPATCH_WIDTH(64) + DISPATCH_WIDTH(128) + DISPATCH_WIDTH(256) +#undef DISPATCH_WIDTH + TORCH_CHECK(false, "dpa1_graph_compress: unsupported width ", width); +} + +template +void dispatch_backward(int width, + long node_count, + long edge_count, + int ntypes, + bool one_side, + bool smooth, + int axis, + bool canonical, + bool masked, + int descriptor_stride, + float rcut, + float rcut_smooth, + float protection, + float inverse_neighbors, + float lower, + float upper, + float table_max, + float stride0, + float stride1, + const torch::Tensor& descriptor_gradient, + const float* rotation_gradient, + const torch::Tensor& moment, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& atype, + const torch::Tensor& average, + const torch::Tensor& inverse_stddev, + const torch::Tensor& table, + const torch::Tensor& gate_table, + torch::Tensor& edge_gradient, + cudaStream_t stream) { +#define DISPATCH_WIDTH(value) \ + if (width == value) { \ + if (canonical && !masked) { \ + launch_backward( \ + node_count, edge_count, ntypes, one_side, smooth, axis, \ + descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, \ + lower, upper, table_max, stride0, stride1, descriptor_gradient, \ + rotation_gradient, moment, edge_vec, edge_index, edge_mask, \ + destination_order, destination_row_ptr, atype, average, \ + inverse_stddev, table, gate_table, edge_gradient, stream); \ + } else if (canonical) { \ + launch_backward( \ + node_count, edge_count, ntypes, one_side, smooth, axis, \ + descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, \ + lower, upper, table_max, stride0, stride1, descriptor_gradient, \ + rotation_gradient, moment, edge_vec, edge_index, edge_mask, \ + destination_order, destination_row_ptr, atype, average, \ + inverse_stddev, table, gate_table, edge_gradient, stream); \ + } else { \ + launch_backward( \ + node_count, edge_count, ntypes, one_side, smooth, axis, \ + descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, \ + lower, upper, table_max, stride0, stride1, descriptor_gradient, \ + rotation_gradient, moment, edge_vec, edge_index, edge_mask, \ + destination_order, destination_row_ptr, atype, average, \ + inverse_stddev, table, gate_table, edge_gradient, stream); \ + } \ + return; \ + } + DISPATCH_WIDTH(8) + DISPATCH_WIDTH(16) + DISPATCH_WIDTH(32) + DISPATCH_WIDTH(64) + DISPATCH_WIDTH(128) + DISPATCH_WIDTH(256) +#undef DISPATCH_WIDTH + TORCH_CHECK(false, "dpa1_graph_compress_backward: unsupported width ", width); +} + +void validate_inputs(const torch::Tensor& edge_vec, + const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& atype, + const torch::Tensor& average, + const torch::Tensor& inverse_stddev, + const torch::Tensor& table, + const torch::Tensor& gate_table, + int width, + int axis) { + TORCH_CHECK(edge_vec.is_cuda() && edge_index.is_cuda() && + edge_mask.is_cuda() && destination_order.is_cuda() && + destination_row_ptr.is_cuda() && atype.is_cuda() && + average.is_cuda() && inverse_stddev.is_cuda() && + table.is_cuda() && gate_table.is_cuda(), + "dpa1_graph_compress: inputs must be CUDA tensors"); + TORCH_CHECK( + edge_vec.is_contiguous() && edge_index.is_contiguous() && + edge_mask.is_contiguous() && destination_order.is_contiguous() && + destination_row_ptr.is_contiguous() && atype.is_contiguous() && + average.is_contiguous() && inverse_stddev.is_contiguous() && + 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::kInt64, + "dpa1_graph_compress: edge_index must be int32 or int64"); + TORCH_CHECK(destination_order.scalar_type() == edge_index.scalar_type(), + "dpa1_graph_compress: destination_order must match the " + "edge_index dtype"); + TORCH_CHECK(edge_mask.scalar_type() == torch::kBool, + "dpa1_graph_compress: edge_mask must be bool"); + TORCH_CHECK(atype.scalar_type() == torch::kInt64, + "dpa1_graph_compress: atype must be int64"); + TORCH_CHECK(destination_row_ptr.scalar_type() == torch::kInt64, + "dpa1_graph_compress: destination_row_ptr must be int64"); + TORCH_CHECK(average.scalar_type() == torch::kFloat32 && + inverse_stddev.scalar_type() == torch::kFloat32 && + table.scalar_type() == torch::kFloat32 && + gate_table.scalar_type() == torch::kFloat32, + "dpa1_graph_compress: statistics and tables must be fp32"); + TORCH_CHECK(axis > 0 && axis <= 16 && axis <= width, + "dpa1_graph_compress: axis must be in [1, min(16, width)]"); +} + +} // namespace + +std::tuple dpa1_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 type_embedding, + torch::Tensor average, + torch::Tensor inverse_stddev, + torch::Tensor table, + torch::Tensor gate_table, + int64_t type_one_side, + int64_t concatenate_type_embedding, + int64_t write_rotation, + int64_t smooth, + int64_t axis, + bool canonical, + double lower, + double upper, + double table_max, + double stride0, + double stride1, + double rcut, + double rcut_smooth, + double protection, + double neighbors) { + const long node_count = atype.size(0); + const int width = static_cast(table.size(1) / 6); + validate_inputs(edge_vec, edge_index, edge_mask, destination_order, + destination_row_ptr, atype, average, inverse_stddev, table, + gate_table, width, static_cast(axis)); + TORCH_CHECK(type_embedding.is_cuda() && type_embedding.is_contiguous() && + type_embedding.scalar_type() == torch::kFloat32, + "dpa1_graph_compress: type_embedding must be contiguous fp32 " + "on CUDA"); + const int ntypes = static_cast(type_embedding.size(0)); + const int type_embedding_dim = static_cast(type_embedding.size(1)); + const int output_dim = width * static_cast(axis) + + (concatenate_type_embedding ? type_embedding_dim : 0); + auto options = edge_vec.options().dtype(torch::kFloat32); + auto descriptor = torch::empty({node_count, output_dim}, options); + auto rotation = + torch::empty({write_rotation ? node_count : 0, width, 3}, options); + auto moment = torch::empty({node_count, 4, width}, options); + if (node_count == 0) { + return {descriptor, rotation, moment}; + } + const auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous(); + const auto stream = at::cuda::getCurrentCUDAStream(); + + auto launch = [&](auto index_tag) { + using index_t = decltype(index_tag); + dispatch_forward( + width, node_count, ntypes, type_one_side != 0, smooth != 0, + static_cast(axis), canonical, edge_mask.numel() != 0, + concatenate_type_embedding != 0, write_rotation != 0, + type_embedding_dim, static_cast(rcut), + static_cast(rcut_smooth), static_cast(protection), + static_cast(1.0 / neighbors), static_cast(lower), + static_cast(upper), static_cast(table_max), + static_cast(stride0), static_cast(stride1), + edge_vec_float, edge_index, edge_mask, destination_order, + destination_row_ptr, atype, type_embedding, average, inverse_stddev, + table, gate_table, descriptor, rotation, moment, stream); + }; + if (edge_index.scalar_type() == torch::kInt32) { + launch(int{}); + } else { + launch(long{}); + } + return {descriptor, rotation, moment}; +} + +torch::Tensor dpa1_graph_compress_backward( + torch::Tensor descriptor_gradient, + std::optional rotation_gradient, + torch::Tensor moment, + 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 average, + torch::Tensor inverse_stddev, + torch::Tensor table, + torch::Tensor gate_table, + int64_t type_one_side, + int64_t smooth, + int64_t axis, + bool canonical, + double lower, + double upper, + double table_max, + double stride0, + double stride1, + double rcut, + double rcut_smooth, + double protection, + double neighbors) { + const long node_count = atype.size(0); + const long edge_count = edge_vec.size(0); + const int width = static_cast(table.size(1) / 6); + validate_inputs(edge_vec, edge_index, edge_mask, destination_order, + destination_row_ptr, atype, average, inverse_stddev, table, + gate_table, width, static_cast(axis)); + if (node_count == 0) { + return torch::zeros_like(edge_vec); + } + const int ntypes = type_one_side + ? static_cast(gate_table.size(0)) + : static_cast(llround( + sqrt(static_cast(gate_table.size(0))))); + auto descriptor_gradient_float = + descriptor_gradient.to(torch::kFloat32).contiguous(); + torch::Tensor rotation_gradient_float; + const float* rotation_gradient_ptr = nullptr; + if (rotation_gradient.has_value() && rotation_gradient->defined() && + rotation_gradient->numel() > 0) { + rotation_gradient_float = + rotation_gradient->to(torch::kFloat32).contiguous(); + rotation_gradient_ptr = rotation_gradient_float.data_ptr(); + } + auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous(); + auto edge_gradient = torch::empty_like(edge_vec_float); + const auto stream = at::cuda::getCurrentCUDAStream(); + + auto launch = [&](auto index_tag) { + using index_t = decltype(index_tag); + dispatch_backward( + width, node_count, edge_count, ntypes, type_one_side != 0, smooth != 0, + static_cast(axis), canonical, edge_mask.numel() != 0, + static_cast(descriptor_gradient_float.size(1)), + static_cast(rcut), static_cast(rcut_smooth), + static_cast(protection), static_cast(1.0 / neighbors), + static_cast(lower), static_cast(upper), + static_cast(table_max), static_cast(stride0), + static_cast(stride1), descriptor_gradient_float, + rotation_gradient_ptr, moment, edge_vec_float, edge_index, edge_mask, + destination_order, destination_row_ptr, atype, average, inverse_stddev, + table, gate_table, edge_gradient, stream); + }; + if (edge_index.scalar_type() == torch::kInt32) { + launch(int{}); + } else { + launch(long{}); + } + return edge_gradient.to(edge_vec.scalar_type()); +} + +std::tuple dpa1_canonical_compress( + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor type_embedding, + torch::Tensor average, + torch::Tensor inverse_stddev, + torch::Tensor table, + torch::Tensor gate_table, + int64_t type_one_side, + int64_t concatenate_type_embedding, + int64_t write_rotation, + int64_t smooth, + int64_t axis, + double lower, + double upper, + double table_max, + double stride0, + double stride1, + double rcut, + double rcut_smooth, + double protection, + double neighbors) { + TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), + "dpa1_canonical_compress: source and edge_vec storage must " + "share the edge axis"); + TORCH_CHECK(destination_row_ptr.numel() == atype.size(0) + 1, + "dpa1_canonical_compress: destination_row_ptr must have N + 1 " + "entries"); + auto edge_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); + auto destination_order = torch::empty({0}, source.options()); + return dpa1_graph_compress( + edge_vec, source, edge_mask, destination_order, destination_row_ptr, + atype, type_embedding, average, inverse_stddev, table, gate_table, + type_one_side, concatenate_type_embedding, write_rotation, smooth, axis, + true, lower, upper, table_max, stride0, stride1, rcut, rcut_smooth, + protection, neighbors); +} + +torch::Tensor dpa1_canonical_compress_backward( + torch::Tensor descriptor_gradient, + std::optional rotation_gradient, + torch::Tensor moment, + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor atype, + torch::Tensor average, + torch::Tensor inverse_stddev, + torch::Tensor table, + torch::Tensor gate_table, + int64_t type_one_side, + int64_t smooth, + int64_t axis, + double lower, + double upper, + double table_max, + double stride0, + double stride1, + double rcut, + double rcut_smooth, + double protection, + double neighbors) { + TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), + "dpa1_canonical_compress_backward: source and edge_vec storage " + "must share the edge axis"); + TORCH_CHECK(destination_row_ptr.numel() == atype.size(0) + 1, + "dpa1_canonical_compress_backward: destination_row_ptr must " + "have N + 1 entries"); + auto edge_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); + auto destination_order = torch::empty({0}, source.options()); + return dpa1_graph_compress_backward( + descriptor_gradient, rotation_gradient, moment, edge_vec, source, + edge_mask, destination_order, destination_row_ptr, atype, average, + inverse_stddev, table, gate_table, type_one_side, smooth, axis, true, + lower, upper, table_max, stride0, stride1, rcut, rcut_smooth, protection, + neighbors); +} + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.def( + "dpa1_graph_compress(Tensor edge_vec, Tensor edge_index, " + "Tensor edge_mask, Tensor destination_order, " + "Tensor destination_row_ptr, Tensor atype, " + "Tensor type_embedding, Tensor average, Tensor inverse_stddev, " + "Tensor table, Tensor gate_table, int type_one_side, " + "int concatenate_type_embedding, int write_rotation, int smooth, " + "int axis, bool canonical, float lower, float upper, float table_max, " + "float stride0, float stride1, " + "float rcut, float rcut_smooth, float protection, float neighbors) " + "-> (Tensor descriptor, Tensor rotation, Tensor moment)"); + library.impl("dpa1_graph_compress", torch::kCUDA, &dpa1_graph_compress); + library.def( + "dpa1_graph_compress_backward(Tensor descriptor_gradient, " + "Tensor? rotation_gradient, Tensor moment, Tensor edge_vec, " + "Tensor edge_index, Tensor edge_mask, Tensor destination_order, " + "Tensor destination_row_ptr, Tensor atype, Tensor average, " + "Tensor inverse_stddev, Tensor table, " + "Tensor gate_table, int type_one_side, int smooth, int axis, " + "bool canonical, float lower, float upper, float table_max, float " + "stride0, " + "float stride1, float rcut, float rcut_smooth, float protection, " + "float neighbors) -> Tensor"); + library.impl("dpa1_graph_compress_backward", torch::kCUDA, + &dpa1_graph_compress_backward); + library.def( + "dpa1_canonical_compress(Tensor edge_vec, Tensor source, " + "Tensor destination_row_ptr, Tensor atype, Tensor type_embedding, " + "Tensor average, Tensor inverse_stddev, Tensor table, " + "Tensor gate_table, int type_one_side, int concatenate_type_embedding, " + "int write_rotation, int smooth, int axis, float lower, float upper, " + "float table_max, float stride0, float stride1, float rcut, " + "float rcut_smooth, float protection, float neighbors) -> " + "(Tensor descriptor, Tensor rotation, Tensor moment)"); + library.impl("dpa1_canonical_compress", torch::kCUDA, + &dpa1_canonical_compress); + library.def( + "dpa1_canonical_compress_backward(Tensor descriptor_gradient, " + "Tensor? rotation_gradient, Tensor moment, Tensor edge_vec, " + "Tensor source, Tensor destination_row_ptr, Tensor atype, " + "Tensor average, Tensor inverse_stddev, Tensor table, " + "Tensor gate_table, int type_one_side, int smooth, int axis, " + "float lower, float upper, float table_max, float stride0, " + "float stride1, float rcut, float rcut_smooth, float protection, " + "float neighbors) -> Tensor"); + library.impl("dpa1_canonical_compress_backward", torch::kCUDA, + &dpa1_canonical_compress_backward); +} diff --git a/source/op/pt/dpa1_graph_compress_tuning.h b/source/op/pt/dpa1_graph_compress_tuning.h new file mode 100644 index 0000000000..304ee79093 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress_tuning.h @@ -0,0 +1,229 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Runtime resource selection for the compressed DPA1 CUDA kernels. + +#pragma once + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace deepmd::dpa1_compress_tuning { + +enum class ResourcePolicy : int { + kBalanced = 0, + kOccupancy = 1, +}; + +enum class KernelDirection : int { + kForward = 0, + kBackward = 1, +}; + +struct LaunchConfig { + ResourcePolicy resource; + int threads; +}; + +struct TuningKey { + int device; + int direction; + int width; + int axis; + int canonical; + int index_bytes; + int model_flags; + int model_stride; + int type_class; + int size_class; + int degree_class; + + bool operator==(const TuningKey& other) const { + return device == other.device && direction == other.direction && + width == other.width && axis == other.axis && + canonical == other.canonical && index_bytes == other.index_bytes && + model_flags == other.model_flags && + model_stride == other.model_stride && + type_class == other.type_class && size_class == other.size_class && + degree_class == other.degree_class; + } +}; + +struct TuningKeyHash { + std::size_t operator()(const TuningKey& key) const { + std::size_t value = 0; + const std::array fields = { + key.device, key.direction, key.width, key.axis, + key.canonical, key.index_bytes, key.model_flags, key.model_stride, + key.type_class, key.size_class, key.degree_class}; + for (const int field : fields) { + value ^= + std::hash{}(field) + 0x9e3779b9 + (value << 6) + (value >> 2); + } + return value; + } +}; + +inline std::mutex& tuning_cache_mutex() { + static std::mutex mutex; + return mutex; +} + +inline std::unordered_map& +tuning_cache() { + static std::unordered_map cache; + return cache; +} + +inline const cudaDeviceProp& device_properties(int device) { + constexpr int kMaximumCachedDevices = 64; + TORCH_CHECK(device >= 0 && device < kMaximumCachedDevices, + "dpa1_graph_compress: unsupported CUDA device index ", device); + static std::array initialized; + static std::array properties; + std::call_once(initialized[device], [device] { + TORCH_CHECK( + cudaGetDeviceProperties(&properties[device], device) == cudaSuccess, + "dpa1_graph_compress: cannot query CUDA device properties"); + }); + return properties[device]; +} + +inline LaunchConfig architecture_fallback(const cudaDeviceProp& properties) { + if (properties.major >= 9) { + if (properties.multiProcessorCount <= 80) { + return {ResourcePolicy::kOccupancy, 256}; + } + return {ResourcePolicy::kBalanced, 256}; + } + if (properties.major == 8) { + return {ResourcePolicy::kBalanced, 256}; + } + if (properties.major == 7) { + return {ResourcePolicy::kBalanced, 128}; + } + return {ResourcePolicy::kBalanced, 256}; +} + +inline int workload_size_class(long node_count, int multiprocessor_count) { + const long nodes_per_sm = node_count / std::max(multiprocessor_count, 1); + if (nodes_per_sm < 8) { + return 0; + } + if (nodes_per_sm < 64) { + return 1; + } + return 2; +} + +inline int workload_degree_class(long node_count, long edge_count) { + const long degree = edge_count / std::max(node_count, 1L); + if (degree < 32) { + return 0; + } + if (degree < 128) { + return 1; + } + return 2; +} + +inline int type_count_class(int type_count) { + if (type_count <= 4) { + return 0; + } + if (type_count <= 16) { + return 1; + } + return 2; +} + +template +LaunchConfig select_launch_config(const TuningKey& key, + const cudaDeviceProp& properties, + long node_count, + cudaStream_t stream, + const LaunchFunction& launch) { + const LaunchConfig fallback = architecture_fallback(properties); + + std::lock_guard lock(tuning_cache_mutex()); + auto& cache = tuning_cache(); + const auto cached = cache.find(key); + if (cached != cache.end()) { + return cached->second; + } + + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + const cudaError_t capture_error = + cudaStreamIsCapturing(stream, &capture_status); + if (capture_error != cudaSuccess || + capture_status != cudaStreamCaptureStatusNone) { + return fallback; + } + if (node_count < 128) { + cache.emplace(key, fallback); + return fallback; + } + + const long sample_node_count = std::min( + node_count, + std::max(4096L, static_cast(properties.multiProcessorCount) * 64)); + constexpr std::array kCandidates = {{ + {ResourcePolicy::kBalanced, 128}, + {ResourcePolicy::kBalanced, 256}, + {ResourcePolicy::kOccupancy, 128}, + {ResourcePolicy::kOccupancy, 256}, + }}; + constexpr int kRepetitions = 2; + + cudaEvent_t start = nullptr; + cudaEvent_t stop = nullptr; + if (cudaEventCreate(&start) != cudaSuccess || + cudaEventCreate(&stop) != cudaSuccess) { + if (start != nullptr) { + cudaEventDestroy(start); + } + if (stop != nullptr) { + cudaEventDestroy(stop); + } + cache.emplace(key, fallback); + return fallback; + } + + LaunchConfig best = fallback; + float best_milliseconds = std::numeric_limits::infinity(); + for (const LaunchConfig& candidate : kCandidates) { + launch(candidate, sample_node_count); + if (cudaPeekAtLastError() != cudaSuccess || + cudaEventRecord(start, stream) != cudaSuccess) { + continue; + } + for (int repetition = 0; repetition < kRepetitions; ++repetition) { + launch(candidate, sample_node_count); + } + if (cudaPeekAtLastError() != cudaSuccess || + cudaEventRecord(stop, stream) != cudaSuccess || + cudaEventSynchronize(stop) != cudaSuccess) { + continue; + } + float milliseconds = 0.0f; + if (cudaEventElapsedTime(&milliseconds, start, stop) == cudaSuccess) { + milliseconds /= kRepetitions; + if (milliseconds < best_milliseconds) { + best_milliseconds = milliseconds; + best = candidate; + } + } + } + cudaEventDestroy(start); + cudaEventDestroy(stop); + cache.emplace(key, best); + return best; +} + +} // namespace deepmd::dpa1_compress_tuning diff --git a/source/op/pt/dpa1_graph_descriptor.cu b/source/op/pt/dpa1_graph_descriptor.cu new file mode 100644 index 0000000000..4bf2d28148 --- /dev/null +++ b/source/op/pt/dpa1_graph_descriptor.cu @@ -0,0 +1,1499 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Fused CUDA descriptor for the DPA1 (``se_atten``) graph lower -- the +// attention-free configuration (concat or strip tebd input) evaluated on +// the flat edge stream. One forward kernel produces the per-node moment +// matrix from ``edge_vec``; one backward kernel returns ``d(edge_vec)`` for +// the analytic force / virial assembly. No ``(E, .)`` activation tensor is +// ever materialized on the autograd tape: the embedding runs on +// shared-memory tiles, and the two operands the backward cannot recompute +// cheaply are spilled once in a streaming layout (see "Saved tensors" +// below). +// +// Mathematical contract (matches ``DescrptBlockSeAtten`` at attn_layer 0): +// rr = ((sw/q, sw*x/q^2, sw*y/q^2, sw*z/q^2) - davg[ct]) / dstd[ct] +// with q = |r| + protection, sw the quintic smooth switch of |r| on +// [rcut_smth, rcut]; masked (padding) edges use |r| + 1. +// in = [rr0, tebd[nt], tebd[ct]] (concat) | rr0 (strip) +// h1 = act(in @ W1 + b1) * idt1 +// h2 = act(h1 @ W2 + b2) * idt2 (+ [h1, h1] doubling residual) +// g = act(h2 @ W3 + b3) * idt3 (+ [h2, h2] doubling residual) +// gg = g (concat) +// = g * (1 + gate[pair] (* sw if smooth)) (strip type-pair gate, +// gate = embeddings_strip(tebd pairs), precomputed by the caller) +// gr[n, i, k] = (1/nnei) * segment_sum(gg[:, i] * mask * rr[:, k], dst) +// grrg[n, i * axis + j] = sum_k gr[n, i, k] * gr[n, j, k] (j < axis) +// rot_mat[n, i, c] = gr[n, i, 1 + c] +// grrg tail (concat_output_tebd): type_embedding[atype[n]]. +// The moment is stored transposed as (N, 4, NG); both contractions above are +// invariant to that layout, so downstream results equal the reference up to +// fp32 summation order. +// +// Algebraic reductions +// -------------------- +// * Pair table: the concat embedding input is [rr0 | tebd[nt] | tebd[ct]], +// whose type-embedding block takes only ``ntypes^2`` distinct values. The +// host folds those rows of W1 (plus b1) into a per-type-pair table of +// layer-1 pre-activations, so layer 1 degenerates to one FMA per channel: +// ``pre1 = rr0 * W1[0] + pair_table[ct * T + nt]``. This removes the +// layer-1 GEMM and every per-edge type-embedding gather. Strip mode keeps +// the same code path with the table collapsed to its single bias row, and +// applies the type-pair gate table in the layer-3 epilogue instead. +// * Sort-free CSR: graph builders emit real edges already dst-sorted, so the +// edge ordering is a histogram + exclusive scan + identity permutation; +// genuinely unsorted input falls back to an atomic-cursor scatter. This +// replaces a per-step argsort. +// +// Kernel organization +// ------------------- +// * Every dense stage is an 8x8 (edges x channels) register micro-tile whose +// activation operand is one float4 shared-memory fragment and whose weight +// operand is a warp-uniform float4 ``__ldg`` broadcast (L1-resident, +// one transaction per warp). Weights are never staged in shared memory; +// the same row-major layout serves the forward k-loops and the dgrad +// loops (which fix the input channel's row and walk the output axis). +// * The moment reduction walks the tile's CSR runs per (moment row, channel) +// slot and flushes one global atomicAdd per run -- no shuffle reductions +// on the forward path. +// * The backward gathers the upstream moment gradient once per (edge, run): +// rows of dgr for the tile's first RMAX runs are staged in shared memory, +// and an all-edges-share-one-run fast path amortizes the four scalars per +// channel over the whole eight-edge fragment. Per-edge partials of the +// environment gradient accumulate in per-warp shared-memory banks (a +// shuffle fold halves the writers), avoiding shared-memory atomics on the +// hot path. +// +// Saved tensors (forward -> backward) +// ----------------------------------- +// ``pre2`` (N2, E_pad) and -- tanh: ``g``; silu: ``pre3`` -- (NG, E_pad) are +// spilled transposed with streaming (evict-first) stores so both sides read +// coalesced rows without evicting the L1-resident weights. This halves the +// backward FLOPs relative to a full recompute at the cost of +// ``(N2 + NG) * 4`` bytes per edge. E_pad rounds E up to the tile size so +// partial-tile vector stores stay in bounds. +// +// Applicability (enforced by the Python gate): three embedding layers with +// N1 in {8, 16, 32, 64}, N2 in {N1, 2*N1}, NG in {N2, 2*N2}, N2 <= 64, and +// NG <= 128. Layers 2 and 3 implement identity and width-doubling residuals. +// Layer 1 is accepted only when its native input/output shape does not form a +// residual. All layers share tanh or silu, with optional timesteps, fp32 +// weights, statistics, and compute. ``edge_vec`` in fp32 or fp64 is cast to +// fp32 on entry and the leaf dtype is restored on the returned gradient. +// ``concat`` or ``strip`` tebd input (strip with or without the smooth gate), +// ``attn_layer == 0``, no excluded type pairs. + +#include +#include +#include + +#include "dpa1_graph_common.cuh" + +namespace { + +// Activation codes follow deepmd.kernels.triton.dpa1.activation.ACT_CODES: +// 0 = tanh, 1 = silu. Forward and backward share this helper so energy and +// its analytic force gradient stay consistent (the potential-energy surface +// remains smooth). +// +// The sigmoid factor of silu(z) = z * sigmoid(z) is evaluated through the +// identity sigmoid(z) = 0.5 * (1 + tanh(0.5 * z)). The accurate fp32 division +// in 1 / (1 + expf(-z)) -- not the exponential -- dominates the silu activation +// cost here; the tanh form removes that division, and since this architecture +// emits no dedicated tanh hardware instruction the tanh factor is no costlier +// than the exponential it replaces. Every silu site -- the forward value, its +// derivative, and the backward reconstruction of the pre-activation -- must use +// this identical expression; a mismatched sigmoid on either side would break +// the identity force == d(energy)/d(coord) at the fp32 level. +template +DEV_INLINE float act_val(float z) { + if constexpr (ACT == 0) { + return tanhf(z); + } + return 0.5f * z * (1.f + tanhf(0.5f * z)); +} + +// Activation value and derivative at z. +template +DEV_INLINE float2 act_vg(float z) { + if constexpr (ACT == 0) { + float a = tanhf(z); + return make_float2(a, 1.f - a * a); + } + const float s = 0.5f * (1.f + tanhf(0.5f * z)); // sigmoid via tanh identity + return make_float2(z * s, s * (1.f + z * (1.f - s))); +} + +// ====================================================================== +// Pair table: fold the type-embedding rows of W1 (and b1) into a per-pair +// layer-1 pre-activation. Two-side pairs index as p = ct * T + nt; one-side +// as p = nt. +// ====================================================================== +__global__ void pair_table_kernel(int n_pairs, + int ntypes, + int tebd_dim, + int n1, + int one_side, + const float* __restrict__ tebd, + const float* __restrict__ w1, + const float* __restrict__ b1, + float* __restrict__ pair_table) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= n_pairs * n1) { + return; + } + const int p = idx / n1, o = idx % n1; + const int nt = one_side ? p : p % ntypes; + const int ct = one_side ? 0 : p / ntypes; + float acc = b1 ? b1[o] : 0.f; + for (int t = 0; t < tebd_dim; ++t) { + acc = fmaf(tebd[nt * tebd_dim + t], w1[(1 + t) * n1 + o], acc); + } + if (!one_side) { + for (int t = 0; t < tebd_dim; ++t) { + acc = fmaf(tebd[ct * tebd_dim + t], w1[(1 + tebd_dim + t) * n1 + o], acc); + } + } + pair_table[(long)p * n1 + o] = acc; +} + +// Layer-1 tile from the pair table: float4 blocks of the (BE, N1) sheet, so +// the gather streams coalesced 16-byte words. Strip mode has no per-pair +// layer-1 term (the type embedding enters through the gate instead), so the +// table degenerates to its single bias row. +template +DEV_INLINE void layer1_tile(int tid, + const EdgeTablesT& T, + int strip, + const float* __restrict__ pair_table, + const float* __restrict__ w1, + const float* __restrict__ idt1, + float (*h1)[STRIDE]) { + constexpr int kBlocks = TILE * N1 / 4; + for (int blk = tid; blk < kBlocks; blk += kThreads) { + const int e = (blk * 4) / N1, o = (blk * 4) % N1; + const long row = strip ? 0 : (long)T.pair_idx[e]; + const float4 pp = + *reinterpret_cast(pair_table + row * N1 + o); + const float4 ws = weight4(w1 + o); + const float4 it = + idt1 ? weight4(idt1 + o) : make_float4(1.f, 1.f, 1.f, 1.f); + const float rv = T.radial[e]; + h1[o + 0][e] = act_val(fmaf(rv, ws.x, pp.x)) * it.x; + h1[o + 1][e] = act_val(fmaf(rv, ws.y, pp.y)) * it.y; + h1[o + 2][e] = act_val(fmaf(rv, ws.z, pp.z)) * it.z; + h1[o + 3][e] = act_val(fmaf(rv, ws.w, pp.w)) * it.w; + } +} + +// ====================================================================== +// Forward kernel. +// +// Persistent CTAs stride over BE-edge slices of the CSR-ordered stream. +// Stages per tile (tx = tid % 16 indexes eight-edge fragments, ty = tid / 16 +// indexes channel groups): +// 0. environment matrix + run scan (stage_tile) +// 1. h1 tile from the pair table (layer1_tile) +// 2. GEMM2 (8e x 4c micro-tiles) -> h2 tile; spill pre2 +// 3. GEMM3 (8e x 8c micro-tiles) -> g tile (overlaying h1 through the +// stage union) + spill g (tanh) or pre3 (silu) +// 4. moment walk: thread owns channel c, streams each CSR run over the +// tile rows with all four moment components, one atomicAdd per +// (run, component, channel) +// ====================================================================== +template +struct FwdSmem { + static constexpr int STRIDE = TILE + 4; + // The g tile overlays the (dead) h1 tile; NG >= N1, so the union is sized + // by g. h1's last readers (the GEMM2 residual epilogue) finish before the + // barrier that precedes the first g write. + union { + float h1[N1][STRIDE]; // stages 1-2 + float g[NG][STRIDE]; // stages 3-4 (walk view) + } u; + float h2[N2][STRIDE]; + EdgeTablesT T; +}; + +static_assert(sizeof(FwdSmem<32, 64, 128, 128>) > 96 * 1024); +static_assert(sizeof(FwdSmem<32, 64, 128, 64>) <= 96 * 1024); + +template +__global__ __launch_bounds__(kThreads, 2) void dpa1_graph_forward_kernel( + long n_edge, + int ntypes, + int one_side, + int resnet2, + int resnet3, + int smooth, + float rcut, + float rcut_smth, + float protection, + float inv_nnei, + const float* __restrict__ edge_vec, + const long* __restrict__ edge_index, + const bool* __restrict__ edge_mask, + const long* __restrict__ atype, + const float* __restrict__ davg, + const float* __restrict__ inv_dstd, + const float* __restrict__ w1, + const float* __restrict__ w2, + const float* __restrict__ b2, + const float* __restrict__ idt1, + const float* __restrict__ idt2, + const float* __restrict__ w3, + const float* __restrict__ b3, + const float* __restrict__ idt3, + const float* __restrict__ pair_table, + const float* __restrict__ gate_table, + const int* __restrict__ order, + float* __restrict__ gr, + float* __restrict__ pre2_saved, + float* __restrict__ g_saved, + long e_pad) { + constexpr int TILE = 16 * EPT; // edges per tile (16 edge-lanes x EPT) + constexpr int STRIDE = TILE + 4; + extern __shared__ char smem_raw[]; + auto& S = *reinterpret_cast*>(smem_raw); + const int tid = threadIdx.x; + const int tx = tid % 16, ty = tid / 16; + const int erow = tx * EPT; + constexpr int strip = STRIP; + + const long n_tiles = ceil_div(n_edge, TILE); + for (long tile = blockIdx.x; tile < n_tiles; tile += gridDim.x) { + const long tile_base = tile * TILE; + const int rows = (int)min((long)TILE, n_edge - tile_base); + + stage_tile(tid, tile_base, rows, n_edge, ntypes, one_side, rcut, + rcut_smth, protection, inv_nnei, edge_vec, edge_index, + edge_mask, atype, davg, inv_dstd, order, S.T); + layer1_tile(tid, S.T, strip, pair_table, w1, idt1, + S.u.h1); + __syncthreads(); + + // === Step 2. GEMM2 -> h2 tile; spill pre2 (transposed (N2, E)) === + // The residual reads h1[c % N1], which covers both the doubling + // ([h1, h1], N2 == 2 * N1) and the identity (N2 == N1) layer shapes. + { + const int c0 = ty * 4; // 4 channels per group; groups beyond N2 idle + if (c0 < N2) { + float acc[4][EPT]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < EPT; ++i) { + acc[j][i] = 0.f; + } +#pragma unroll 8 + for (int k = 0; k < N1; ++k) { + float a[EPT]; + loadN(&S.u.h1[k][erow], a); + const float4 b = weight4(w2 + k * N2 + c0); +#pragma unroll + for (int i = 0; i < EPT; ++i) { + acc[0][i] = fmaf(a[i], b.x, acc[0][i]); + acc[1][i] = fmaf(a[i], b.y, acc[1][i]); + acc[2][i] = fmaf(a[i], b.z, acc[2][i]); + acc[3][i] = fmaf(a[i], b.w, acc[3][i]); + } + } + const float4 bias = b2 ? weight4(b2 + c0) : make_float4(0, 0, 0, 0); + const float4 idt = idt2 ? weight4(idt2 + c0) : make_float4(1, 1, 1, 1); + const float bs[4] = {bias.x, bias.y, bias.z, bias.w}; + const float is[4] = {idt.x, idt.y, idt.z, idt.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int c = c0 + j; + float pre[EPT], out[EPT]; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + pre[i] = acc[j][i] + bs[j]; + float v = act_val(pre[i]) * is[j]; + if (resnet2) { + v += S.u.h1[c % N1][erow + i]; + } + out[i] = v; + } + storeN(&S.h2[c][erow], out); + storeN_streaming(&pre2_saved[(long)c * e_pad + tile_base + erow], + pre); + } + } + } + __syncthreads(); + + // === Step 3. GEMM3 -> g tile (over h1's storage) + spill === + // h1's last readers (the GEMM2 residual epilogue) finished at the + // preceding barrier, so the union overlay is race-free. The residual + // reads h2[c % N2] (doubling or identity, as in Step 2). + { + const int c0 = ty * 8; // 8 channels per group; groups beyond NG idle + if (c0 < NG) { + float acc[8][EPT]; // [8 channels][EPT edges] +#pragma unroll + for (int j = 0; j < 8; ++j) +#pragma unroll + for (int i = 0; i < EPT; ++i) { + acc[j][i] = 0.f; + } +#pragma unroll 4 + for (int k = 0; k < N2; ++k) { + float a[EPT], b[8]; // a: EPT edges; b: 8 channel weights + loadN(&S.h2[k][erow], a); + const float4 b0 = weight4(w3 + k * NG + c0); + const float4 b1 = weight4(w3 + k * NG + c0 + 4); + b[0] = b0.x; + b[1] = b0.y; + b[2] = b0.z; + b[3] = b0.w; + b[4] = b1.x; + b[5] = b1.y; + b[6] = b1.z; + b[7] = b1.w; +#pragma unroll + for (int j = 0; j < 8; ++j) // 8 channels +#pragma unroll + for (int i = 0; i < EPT; ++i) { // EPT edges + acc[j][i] = fmaf(a[i], b[j], acc[j][i]); + } + } +#pragma unroll + for (int j = 0; j < 8; ++j) { // 8 channels + const int c = c0 + j; + const float bs = b3 ? __ldg(b3 + c) : 0.f; + const float is = idt3 ? __ldg(idt3 + c) : 1.f; + float out[EPT], pre[EPT]; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + pre[i] = acc[j][i] + bs; + float v = act_val(pre[i]) * is; + if (resnet3) { + v += S.h2[c % N2][erow + i]; + } + out[i] = v; + } + // tanh recovers act' from the saved output (1 - a^2); silu needs + // the pre-activation, so the spill slot holds pre3 instead. The + // spill precedes the strip gate: the backward reconstructs the + // ungated g and regathers the gate. + if constexpr (ACT == 0) { + storeN_streaming(&g_saved[(long)c * e_pad + tile_base + erow], + out); + } else { + storeN_streaming(&g_saved[(long)c * e_pad + tile_base + erow], + pre); + } + // Strip mode: the walk consumes the gated gg = g * (1 + gate_eff). + if (strip) { +#pragma unroll + for (int i = 0; i < EPT; ++i) { + const int e = erow + i; + out[i] *= 1.f + gate_factor(gate_table, NG, S.T.pair_idx[e], c, + S.T.sw[e], smooth); + } + } + storeN(&S.u.g[c][erow], out); + } + } + } + __syncthreads(); + + // === Step 4. Moment walk over CSR runs === + // kThreads / NG threads share one channel and split each run's row + // range; one g read feeds all four moment components. + { + constexpr int kSlices = kThreads / NG; // 2 (NG = 128) or 4 (NG = 64) + const int c = tid % NG; + const int slice = tid / NG; + const float* gcol = &S.u.g[c][0]; + const int n_runs = S.T.n_runs; + for (int run = 0; run < n_runs; ++run) { + const int node = S.T.run_node[run]; + const int rb = S.T.run_begin[run], re = S.T.run_begin[run + 1]; + const int span = re - rb; + const int beg = rb + span * slice / kSlices; + const int end = rb + span * (slice + 1) / kSlices; + float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f; + for (int r = beg; r < end; ++r) { + const float4 rv = *reinterpret_cast(&S.T.rr[r][0]); + const float gv = gcol[r]; + a0 = fmaf(rv.x, gv, a0); + a1 = fmaf(rv.y, gv, a1); + a2 = fmaf(rv.z, gv, a2); + a3 = fmaf(rv.w, gv, a3); + } + if (node >= 0) { + const long base = ((long)node * 4) * NG + c; + if (a0 != 0.f) { + atomicAdd(&gr[base + 0 * NG], a0); + } + if (a1 != 0.f) { + atomicAdd(&gr[base + 1 * NG], a1); + } + if (a2 != 0.f) { + atomicAdd(&gr[base + 2 * NG], a2); + } + if (a3 != 0.f) { + atomicAdd(&gr[base + 3 * NG], a3); + } + } + } + } + __syncthreads(); + } +} + +// ====================================================================== +// Backward kernel. +// +// Consumes the saved pre2 / g spills instead of recomputing the embedding. +// Stages per tile: +// 0. environment matrix + run scan; stage the tile's first RMAX distinct +// dgr rows in shared memory +// 1. h2 tile reconstructed as act(pre2) * idt2 (+ residual re-derived from +// the pair table) +// 2. NG / N2 passes over N2-channel blocks: per (4c x 8e) micro-tile read +// the saved g, form +// dgv = sum_k rr[k] * dgr[dst, k, c] (moment backward) +// dpre3 = dgs * idt3 * act'(pre3) -> block tile x_t +// drr[k] += gg * dgr[dst, k, c] (environment backward) +// dh2[c mod N2] += dgs[c] (layer-3 residual fold) +// with dgs the (strip-gated) g gradient. The N2-wide blocking makes +// the fold thread-local for BOTH residual shapes: block b covers +// channels [b * N2, (b + 1) * N2), so c mod N2 always lands on the +// owning thread's dh2 registers (doubling folds dgs[c] + dgs[c + N2] +// over two blocks; identity folds dgs[c] over the single block). Then +// the dgrad3 block-pass dh2 += dpre3[block] @ W3[:, block]^T. When all +// eight edges of a fragment share one CSR run (the common case) the +// four dgr scalars per channel are read once and amortized across the +// fragment. +// 3. dpre2 = dh2 * act'(pre2) * idt2 -> x_t rows [0, N2); raw dh2 parks in +// the (dead) h2 tile for the layer-2 residual fold +// 4. dgrad2 -> dh1; dpre1 through the pair table; d(radial) reduced into +// shared memory +// 5. analytic environment backward -> d_edge_vec +// ====================================================================== +template +struct BwdSmem { + static constexpr int kResidentRuns = 4; + static constexpr int STRIDE = TILE + 4; + float h2[N2][STRIDE]; // act(pre2) tile; raw dh2 after stage 3 + float x_t[N2][STRIDE]; // dpre3 block tile; dpre2 rows [0, N2) + float drr_banks[8][4][TILE]; // per-warp env-gradient banks + float d_radial[TILE]; + float d_sw[TILE]; // strip gate: dE/d(sw) through gate * sw + float dgr_rows[kResidentRuns][4 * NG]; + int run_slot[TILE]; // resident dgr slot per row (-1: global read) + EdgeTablesT T; +}; + +template +__global__ __launch_bounds__(kThreads, 2) void dpa1_graph_backward_kernel( + long n_edge, + int ntypes, + int one_side, + int resnet2, + int resnet3, + int smooth, + float rcut, + float rcut_smth, + float protection, + float inv_nnei, + const float* __restrict__ edge_vec, + const long* __restrict__ edge_index, + const bool* __restrict__ edge_mask, + const long* __restrict__ atype, + const float* __restrict__ davg, + const float* __restrict__ inv_dstd, + const float* __restrict__ w1, + const float* __restrict__ w2, + const float* __restrict__ b2, + const float* __restrict__ idt1, + const float* __restrict__ idt2, + const float* __restrict__ w3, + const float* __restrict__ b3, + const float* __restrict__ idt3, + const float* __restrict__ pair_table, + const float* __restrict__ gate_table, + const int* __restrict__ order, + const float* __restrict__ dgr, + const float* __restrict__ pre2_saved, + const float* __restrict__ g_saved, + float* __restrict__ d_edge_vec, + long e_pad) { + constexpr int kBlocks = NG / N2; // channel blocks: 1 (identity) or 2 + constexpr int TILE = 16 * EPT; // edges per tile (16 edge-lanes x EPT) + constexpr int kResidentRuns = BwdSmem::kResidentRuns; + extern __shared__ char smem_raw[]; + auto& S = *reinterpret_cast*>(smem_raw); + const int tid = threadIdx.x; + const int tx = tid % 16, ty = tid / 16; + const int erow = tx * EPT; + constexpr int strip = STRIP; + + const long n_tiles = ceil_div(n_edge, TILE); + for (long tile = blockIdx.x; tile < n_tiles; tile += gridDim.x) { + const long tile_base = tile * TILE; + const int rows = (int)min((long)TILE, n_edge - tile_base); + + stage_tile(tid, tile_base, rows, n_edge, ntypes, one_side, rcut, + rcut_smth, protection, inv_nnei, edge_vec, edge_index, + edge_mask, atype, davg, inv_dstd, order, S.T); + if (tid < TILE) { + S.run_slot[tid] = S.T.run_of[tid] < kResidentRuns ? S.T.run_of[tid] : -1; + S.d_radial[tid] = 0.f; + S.d_sw[tid] = 0.f; + } + for (int t = tid; t < 8 * 4 * TILE; t += kThreads) { + (&S.drr_banks[0][0][0])[t] = 0.f; + } + { + const int resident = min(S.T.n_runs, kResidentRuns); + for (int s = 0; s < resident; ++s) { + const long d = S.T.run_node[s]; + if (d < 0) { + continue; + } + for (int t = tid; t < 4 * NG; t += kThreads) { + S.dgr_rows[s][t] = __ldg(dgr + d * 4 * NG + t); + } + } + } + // The staged dgr rows, run slots and cleared accumulator banks are + // written cooperatively across warps; the moment stage below reads them + // from arbitrary warps. + __syncthreads(); + + // === Step 1. h2 tile from the saved pre2 === + // The layer-2 doubling residual re-derives h1 per (channel, edge) from + // the pair table (one FMA and one activation), which is cheaper than a + // second saved tensor. + { + const int c0 = ty * 4; + if (c0 < N2) { + const float4 idt = idt2 ? weight4(idt2 + c0) : make_float4(1, 1, 1, 1); + const float is[4] = {idt.x, idt.y, idt.z, idt.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int c = c0 + j; + float pre[EPT], hv[EPT]; + loadN_streaming(&pre2_saved[(long)c * e_pad + tile_base + erow], + pre); +#pragma unroll + for (int i = 0; i < EPT; ++i) { + hv[i] = act_val(pre[i]) * is[j]; + } + if (resnet2) { + const int cm = c % N1; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + const int e = erow + i; + const long row = strip ? 0 : (long)S.T.pair_idx[e]; + const float it1 = idt1 ? __ldg(idt1 + cm) : 1.f; + const float pre1 = fmaf(S.T.radial[e], __ldg(w1 + cm), + __ldg(pair_table + row * N1 + cm)); + hv[i] += act_val(pre1) * it1; + } + } + storeN(&S.h2[c][erow], hv); + } + } + } + + // Per-edge moment weights, hoisted to registers once per tile (the dgv / + // drr / residual-fold loops otherwise re-read the float4 per channel). + float rr_frag[EPT][4]; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + const float4 rv = *reinterpret_cast(&S.T.rr[erow + i][0]); + rr_frag[i][0] = rv.x; + rr_frag[i][1] = rv.y; + rr_frag[i][2] = rv.z; + rr_frag[i][3] = rv.w; + } +#define RR_(i, k) (rr_frag[i][k]) + + // === Step 2. N2-channel blocks: moment backward + dgrad3 === + // dh2 (4 channels x EPT edges in registers) accumulates the layer-3 + // residual fold and the dgrad3 GEMM across the blocks. + float dh2[4][EPT]; + const int kh0 = ty * 4; +#pragma unroll + for (int block = 0; block < kBlocks; ++block) { + if (block == 0) { +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < EPT; ++i) { + dh2[j][i] = 0.f; + } + } else { + __syncthreads(); // x_t consumed by the previous dgrad3 + } + { + const int c0 = ty * 4; // 4 channels per group; groups beyond N2 idle + if (c0 < N2) { + const int cg = block * N2 + c0; + const int slot0 = S.run_slot[erow]; + const bool uniform_run = + S.T.run_of[erow] == S.T.run_of[erow + EPT - 1] && slot0 >= 0; + float dr0[EPT], dr1[EPT], dr2[EPT], dr3[EPT], dsw[EPT]; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + dr0[i] = dr1[i] = dr2[i] = dr3[i] = dsw[i] = 0.f; + } + // Optional prefetch (PIPE): issue the group's four channel fragments + // of the g / pre3 spill together so their (L2-missing) global + // latency overlaps, then consume from registers. The narrow EPT + // fragment leaves the register headroom this needs. Beneficial for + // wide NG (large spill, e.g. the doubling stacks); for narrow NG it + // adds registers without hiding enough latency, so it is dispatched + // per shape. + float gv4[4][EPT]; + if constexpr (PIPE) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + loadN_streaming( + &g_saved[(long)(cg + j) * e_pad + tile_base + erow], gv4[j]); + } + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int c = cg + j; + const float is = idt3 ? __ldg(idt3 + c) : 1.f; + const float inv_is = 1.f / is; + float gv[EPT], dp[EPT], dgv[EPT]; + if constexpr (PIPE) { +#pragma unroll + for (int i = 0; i < EPT; ++i) { + gv[i] = gv4[j][i]; + } + } else { + loadN_streaming(&g_saved[(long)c * e_pad + tile_base + erow], + gv); + } + if (uniform_run) { + const float* row = S.dgr_rows[slot0]; + const float d0 = row[0 * NG + c], d1 = row[1 * NG + c]; + const float d2 = row[2 * NG + c], d3 = row[3 * NG + c]; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + dgv[i] = RR_(i, 0) * d0 + RR_(i, 1) * d1 + RR_(i, 2) * d2 + + RR_(i, 3) * d3; + } +#pragma unroll + for (int i = 0; i < EPT; ++i) { + float gval, aprime; + if constexpr (ACT == 0) { + gval = gv[i]; + float raw = gval; + if (resnet3) { + raw -= S.h2[c % N2][erow + i]; + } + const float a = raw * inv_is; + aprime = 1.f - a * a; + } else { + const float z = gv[i]; // silu spill holds pre3 + const float s = + 0.5f * + (1.f + + tanhf(0.5f * z)); // sigmoid via tanh; matches act_val + aprime = s * (1.f + z * (1.f - s)); + gval = z * s * is; + if (resnet3) { + gval += S.h2[c % N2][erow + i]; + } + } + // Strip gate chain: gg = g * (1 + gate_eff) feeds the moment, + // so drr uses gg, the g-path gradient scales by (1 + gate_eff) + // and (smooth) the raw gate contributes dE/d(sw). + float gg = gval; + if (strip) { + const int e = erow + i; + const float gate = + __ldg(gate_table + (long)S.T.pair_idx[e] * NG + c); + const float geff = smooth ? gate * S.T.sw[e] : gate; + gg = gval * (1.f + geff); + if (smooth) { + dsw[i] = fmaf(dgv[i] * gval, gate, dsw[i]); + } + dgv[i] *= 1.f + geff; + } + dp[i] = dgv[i] * is * aprime; + dr0[i] = fmaf(gg, d0, dr0[i]); + dr1[i] = fmaf(gg, d1, dr1[i]); + dr2[i] = fmaf(gg, d2, dr2[i]); + dr3[i] = fmaf(gg, d3, dr3[i]); + } + } else { +#pragma unroll + for (int i = 0; i < EPT; ++i) { + const int e = erow + i; + const int slot = S.run_slot[e]; + const float* row = + slot >= 0 ? S.dgr_rows[slot] + : dgr + (long)max(S.T.dst[e], 0) * 4 * NG; + const float d0 = row[0 * NG + c], d1 = row[1 * NG + c]; + const float d2 = row[2 * NG + c], d3 = row[3 * NG + c]; + dgv[i] = RR_(i, 0) * d0 + RR_(i, 1) * d1 + RR_(i, 2) * d2 + + RR_(i, 3) * d3; + float gval, aprime; + if constexpr (ACT == 0) { + gval = gv[i]; + float raw = gval; + if (resnet3) { + raw -= S.h2[c % N2][e]; + } + const float a = raw * inv_is; + aprime = 1.f - a * a; + } else { + const float z = gv[i]; + const float s = + 0.5f * + (1.f + + tanhf(0.5f * z)); // sigmoid via tanh; matches act_val + aprime = s * (1.f + z * (1.f - s)); + gval = z * s * is; + if (resnet3) { + gval += S.h2[c % N2][e]; + } + } + float gg = gval; + if (strip) { + const float gate = + __ldg(gate_table + (long)S.T.pair_idx[e] * NG + c); + const float geff = smooth ? gate * S.T.sw[e] : gate; + gg = gval * (1.f + geff); + if (smooth) { + dsw[i] = fmaf(dgv[i] * gval, gate, dsw[i]); + } + dgv[i] *= 1.f + geff; + } + dp[i] = dgv[i] * is * aprime; + dr0[i] = fmaf(gg, d0, dr0[i]); + dr1[i] = fmaf(gg, d1, dr1[i]); + dr2[i] = fmaf(gg, d2, dr2[i]); + dr3[i] = fmaf(gg, d3, dr3[i]); + } + } + // Layer-3 residual fold: dh2[c mod N2] accumulates dgs over the + // blocks. This thread's fold channels coincide with its dgv + // channels (kh0 == c0), so block b contributes dgs[b * N2 + c] + // -- the doubling shape folds two terms, identity folds one. + if (resnet3) { +#pragma unroll + for (int i = 0; i < EPT; ++i) { + dh2[j][i] += dgv[i]; + } + } + storeN(&S.x_t[c0 + j][erow], dp); + } + // Environment-gradient partials: fold the two channel-group + // halves of each warp (lanes l and l + 16 share the edge + // fragment), then accumulate race-free into this warp's bank. +#pragma unroll + for (int i = 0; i < EPT; ++i) { + dr0[i] += __shfl_xor_sync(0xffffffffu, dr0[i], 16); + dr1[i] += __shfl_xor_sync(0xffffffffu, dr1[i], 16); + dr2[i] += __shfl_xor_sync(0xffffffffu, dr2[i], 16); + dr3[i] += __shfl_xor_sync(0xffffffffu, dr3[i], 16); + dsw[i] += __shfl_xor_sync(0xffffffffu, dsw[i], 16); + } + if ((ty & 1) == 0) { + const int w = tid >> 5; + float* b0 = &S.drr_banks[w][0][erow]; + float* b1 = &S.drr_banks[w][1][erow]; + float* b2v = &S.drr_banks[w][2][erow]; + float* b3v = &S.drr_banks[w][3][erow]; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + b0[i] += dr0[i]; + b1[i] += dr1[i]; + b2v[i] += dr2[i]; + b3v[i] += dr3[i]; + } + if (strip && smooth) { +#pragma unroll + for (int i = 0; i < EPT; ++i) { + if (dsw[i] != 0.f) { + atomicAdd(&S.d_sw[erow + i], dsw[i]); + } + } + } + } + } + } + __syncthreads(); + + // dgrad3 block-pass: dh2 += dpre3[block] @ W3[:, block]^T. + // Four-channel steps keep the weight reads as row-contiguous float4 + // broadcasts. + if (kh0 < N2) { +#pragma unroll 2 + for (int c = 0; c < N2; c += 4) { + float a0[EPT], a1[EPT], a2[EPT], a3[EPT]; + loadN(&S.x_t[c + 0][erow], a0); + loadN(&S.x_t[c + 1][erow], a1); + loadN(&S.x_t[c + 2][erow], a2); + loadN(&S.x_t[c + 3][erow], a3); + const int cg = block * N2 + c; + const float4 w0 = weight4(w3 + (kh0 + 0) * NG + cg); + const float4 w1v = weight4(w3 + (kh0 + 1) * NG + cg); + const float4 w2v = weight4(w3 + (kh0 + 2) * NG + cg); + const float4 w3v = weight4(w3 + (kh0 + 3) * NG + cg); +#pragma unroll + for (int i = 0; i < EPT; ++i) { + dh2[0][i] += + a0[i] * w0.x + a1[i] * w0.y + a2[i] * w0.z + a3[i] * w0.w; + dh2[1][i] += + a0[i] * w1v.x + a1[i] * w1v.y + a2[i] * w1v.z + a3[i] * w1v.w; + dh2[2][i] += + a0[i] * w2v.x + a1[i] * w2v.y + a2[i] * w2v.z + a3[i] * w2v.w; + dh2[3][i] += + a0[i] * w3v.x + a1[i] * w3v.y + a2[i] * w3v.z + a3[i] * w3v.w; + } + } + } + } + + // === Step 3. dpre2 -> x_t rows [0, N2); raw dh2 -> h2 rows === + // h2's activation values are dead once every block consumed them, so + // the tile stores the raw dh2 for the layer-2 residual fold. + __syncthreads(); + if (kh0 < N2) { + const float4 idt = idt2 ? weight4(idt2 + kh0) : make_float4(1, 1, 1, 1); + const float is[4] = {idt.x, idt.y, idt.z, idt.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int k = kh0 + j; + float pre[EPT], dp[EPT]; + loadN_streaming(&pre2_saved[(long)k * e_pad + tile_base + erow], + pre); +#pragma unroll + for (int i = 0; i < EPT; ++i) { + dp[i] = dh2[j][i] * act_vg(pre[i]).y * is[j]; + } + storeN(&S.x_t[k][erow], dp); + storeN(&S.h2[k][erow], dh2[j]); + } + } + __syncthreads(); + + // === Step 4. dgrad2 -> dh1 -> dpre1 -> d(radial) === + // Each channel group owns kC1 = max(N1 / 16, 2) consecutive dh1 + // channels so the 16 groups cover any N1 up to 64. + { + constexpr int kC1 = N1 / 16 > 2 ? N1 / 16 : 2; + const int c0 = ty * kC1; + if (c0 < N1) { + float acc[kC1][EPT]; +#pragma unroll + for (int j = 0; j < kC1; ++j) +#pragma unroll + for (int i = 0; i < EPT; ++i) { + acc[j][i] = 0.f; + } +#pragma unroll 2 + for (int k = 0; k < N2; k += 4) { + float a0[EPT], a1[EPT], a2[EPT], a3[EPT]; + loadN(&S.x_t[k + 0][erow], a0); + loadN(&S.x_t[k + 1][erow], a1); + loadN(&S.x_t[k + 2][erow], a2); + loadN(&S.x_t[k + 3][erow], a3); +#pragma unroll + for (int j = 0; j < kC1; ++j) { + const float4 wv = weight4(w2 + (c0 + j) * N2 + k); +#pragma unroll + for (int i = 0; i < EPT; ++i) { + acc[j][i] += + a0[i] * wv.x + a1[i] * wv.y + a2[i] * wv.z + a3[i] * wv.w; + } + } + } + float ds[EPT]; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + ds[i] = 0.f; + } +#pragma unroll + for (int j = 0; j < kC1; ++j) { + const int c = c0 + j; + const float ws = __ldg(w1 + c); + const float it1 = idt1 ? __ldg(idt1 + c) : 1.f; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + const int e = erow + i; + float dh1 = acc[j][i]; + // Layer-2 residual fold over the raw dh2 parked in the h2 tile: + // one term for the identity shape, two for the doubling shape. + if (resnet2) { +#pragma unroll + for (int b = 0; b < N2 / N1; ++b) { + dh1 += S.h2[c + b * N1][e]; + } + } + const long row = strip ? 0 : (long)S.T.pair_idx[e]; + const float pre1 = + fmaf(S.T.radial[e], ws, __ldg(pair_table + row * N1 + c)); + const float a1d = act_vg(pre1).y * it1; + ds[i] = fmaf(dh1 * a1d, ws, ds[i]); + } + } +#pragma unroll + for (int i = 0; i < EPT; ++i) { + atomicAdd(&S.d_radial[erow + i], ds[i]); + } + } + } + __syncthreads(); + + // === Step 5. Analytic environment backward -> d_edge_vec === + // With v = (sw/q, sw*x/q^2, ...), the chain through sw(|r|), 1/q and the + // per-type normalization is evaluated in closed form. The strip gate + // contributes an additional dE/d(sw) term along the unit vector. + if (tid < rows) { + const int e = order[tile_base + tid]; + if (!edge_mask[e]) { + d_edge_vec[(long)e * 3 + 0] = 0.f; + d_edge_vec[(long)e * 3 + 1] = 0.f; + d_edge_vec[(long)e * 3 + 2] = 0.f; + } else { + const float x = edge_vec[(long)e * 3 + 0]; + const float y = edge_vec[(long)e * 3 + 1]; + const float z = edge_vec[(long)e * 3 + 2]; + const float len = sqrtf(x * x + y * y + z * z); + const float q = len + protection; + const float sw = switch_val(len, rcut_smth, rcut); + const float dsw = switch_deriv(len, rcut_smth, rcut); + const float* isd = inv_dstd + (long)atype[S.T.dst[tid]] * 4; + float dr[4] = {0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int w = 0; w < 8; ++w) +#pragma unroll + for (int k = 0; k < 4; ++k) { + dr[k] += S.drr_banks[w][k][tid]; + } + const float g0 = (dr[0] * inv_nnei + S.d_radial[tid]) * isd[0]; + const float gx = dr[1] * inv_nnei * isd[1]; + const float gy = dr[2] * inv_nnei * isd[2]; + const float gz = dr[3] * inv_nnei * isd[3]; + const float inv_len = len > 0.f ? 1.f / len : 0.f; + const float rq = 1.f / q; + const float gdot = gx * x + gy * y + gz * z; + const float coef = + (g0 * rq * (dsw - sw * rq) + + gdot * rq * rq * (dsw - 2.f * sw * rq) + S.d_sw[tid] * dsw) * + inv_len; + const float s2 = sw * rq * rq; + d_edge_vec[(long)e * 3 + 0] = coef * x + s2 * gx; + d_edge_vec[(long)e * 3 + 1] = coef * y + s2 * gy; + d_edge_vec[(long)e * 3 + 2] = coef * z + s2 * gz; + } + } + __syncthreads(); + } +#undef RR_ +} + +// ====================================================================== +// Host-side dispatch +// ====================================================================== +struct EmbeddingWidths { + int n1, n2, ng; +}; + +// Supported width stacks: each layer keeps (identity residual shape) or +// doubles (concat [x, x] residual shape) the previous width, with N1 a power +// of two in [8, 64], N2 <= 64 and NG <= 128 (the bounds of the 16-group tile +// framework at 256 threads). +EmbeddingWidths check_widths(const torch::Tensor& w1, + const torch::Tensor& w2, + const torch::Tensor& w3) { + EmbeddingWidths s{(int)w1.size(1), (int)w2.size(1), (int)w3.size(1)}; + TORCH_CHECK((int)w2.size(0) == s.n1 && (int)w3.size(0) == s.n2, + "dpa1_graph_descriptor: inconsistent embedding widths"); + TORCH_CHECK( + (s.n2 == s.n1 || s.n2 == 2 * s.n1) && (s.ng == s.n2 || s.ng == 2 * s.n2), + "dpa1_graph_descriptor: each layer width must equal or double the " + "previous one"); + TORCH_CHECK((s.n1 == 8 || s.n1 == 16 || s.n1 == 32 || s.n1 == 64) && + s.n2 <= 64 && s.ng <= 128, + "dpa1_graph_descriptor: unsupported widths (", s.n1, ", ", s.n2, + ", ", s.ng, ")"); + return s; +} + +// CSR ordering + layer-1 pair table, shared by the forward entry point (the +// backward receives both as saved tensors). Strip mode has no type term in +// layer 1: the table degenerates to the single bias row (zeros without a +// bias) and the type embedding enters through the gate table instead. +std::tuple build_order_and_pair_table( + const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + long n_edge, + long n_node, + int ntypes, + int tebd_dim, + int n1, + int one_side, + int strip, + const torch::Tensor& type_embedding, + const torch::Tensor& w1, + const torch::Tensor& b1, + cudaStream_t stream) { + auto order = build_edge_order(edge_index, edge_mask, n_edge, n_node, stream); + + if (strip) { + auto pair_table = b1.numel() + ? b1.reshape({1, n1}).contiguous() + : torch::zeros({1, n1}, type_embedding.options()); + return {order, pair_table}; + } + const long n_pairs = one_side ? ntypes : (long)ntypes * ntypes; + auto pair_table = torch::empty({n_pairs, n1}, type_embedding.options()); + pair_table_kernel<<>>( + (int)n_pairs, ntypes, tebd_dim, n1, one_side, + type_embedding.data_ptr(), w1.data_ptr(), optional_ptr(b1), + pair_table.data_ptr()); + DPA1_CHECK_LAUNCH("dpa1_graph_descriptor pair table"); + return {order, pair_table}; +} + +// Bundles the tensor arguments shared verbatim by the forward and backward +// launches, so the width / activation dispatch stays a one-line macro. +struct LaunchArgs { + long n_edge; + int ntypes, one_side, resnet2, resnet3, smooth; + float rcut, rcut_smth, protection, inv_nnei; + const torch::Tensor &edge_vec, &edge_index, &edge_mask, &atype; + const torch::Tensor &davg, &inv_dstd; + const torch::Tensor &w1, &w2, &b2, &idt1, &idt2, &w3, &b3, &idt3; + const torch::Tensor &pair_table, &gate_table, ℴ + cudaStream_t stream; +}; + +template +void launch_forward(const LaunchArgs& a, + torch::Tensor& gr, + torch::Tensor& pre2_saved, + torch::Tensor& g_saved) { + constexpr int TILE = 16 * EPT; + auto kernel = dpa1_graph_forward_kernel; + const size_t smem = sizeof(FwdSmem); + const cudaError_t attribute_error = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); + TORCH_CHECK( + attribute_error == cudaSuccess, + "dpa1_graph_descriptor forward: cannot configure ", smem, + " bytes of dynamic shared memory: ", cudaGetErrorString(attribute_error)); + kernel<<>>( + a.n_edge, a.ntypes, a.one_side, a.resnet2, a.resnet3, a.smooth, a.rcut, + a.rcut_smth, a.protection, a.inv_nnei, a.edge_vec.data_ptr(), + a.edge_index.data_ptr(), a.edge_mask.data_ptr(), + a.atype.data_ptr(), a.davg.data_ptr(), + a.inv_dstd.data_ptr(), a.w1.data_ptr(), + a.w2.data_ptr(), optional_ptr(a.b2), optional_ptr(a.idt1), + optional_ptr(a.idt2), a.w3.data_ptr(), optional_ptr(a.b3), + optional_ptr(a.idt3), a.pair_table.data_ptr(), + optional_ptr(a.gate_table), a.order.data_ptr(), gr.data_ptr(), + pre2_saved.data_ptr(), g_saved.data_ptr(), + pre2_saved.size(1)); + DPA1_CHECK_LAUNCH("dpa1_graph_descriptor forward"); +} + +template +void launch_forward_portable(const LaunchArgs& a, + torch::Tensor& gr, + torch::Tensor& pre2_saved, + torch::Tensor& g_saved) { + constexpr int kWideEdgesPerThread = 8; + constexpr int kNarrowEdgesPerThread = 4; + constexpr int kWideTile = 16 * kWideEdgesPerThread; + constexpr size_t kWideSharedMemory = sizeof(FwdSmem); + constexpr size_t kPortableSharedMemoryFloor = 48 * 1024; + if constexpr (kWideSharedMemory <= kPortableSharedMemoryFloor) { + launch_forward( + a, gr, pre2_saved, g_saved); + return; + } + const auto* properties = at::cuda::getCurrentDeviceProperties(); + const size_t device_limit = std::max(properties->sharedMemPerBlock, + properties->sharedMemPerBlockOptin); + if (kWideSharedMemory <= device_limit) { + launch_forward( + a, gr, pre2_saved, g_saved); + } else { + constexpr int kNarrowTile = 16 * kNarrowEdgesPerThread; + constexpr size_t kNarrowSharedMemory = + sizeof(FwdSmem); + TORCH_CHECK(kNarrowSharedMemory <= device_limit, + "dpa1_graph_descriptor forward requires ", kNarrowSharedMemory, + " bytes of dynamic shared memory, but the device supports ", + device_limit); + launch_forward( + a, gr, pre2_saved, g_saved); + } +} + +template +void launch_backward(const LaunchArgs& a, + const torch::Tensor& dgr, + const torch::Tensor& pre2_saved, + const torch::Tensor& g_saved, + torch::Tensor& d_edge_vec) { + constexpr int TILE = 16 * EPT; + auto kernel = dpa1_graph_backward_kernel; + const size_t smem = sizeof(BwdSmem); + const cudaError_t attribute_error = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); + TORCH_CHECK( + attribute_error == cudaSuccess, + "dpa1_graph_descriptor backward: cannot configure ", smem, + " bytes of dynamic shared memory: ", cudaGetErrorString(attribute_error)); + kernel<<>>( + a.n_edge, a.ntypes, a.one_side, a.resnet2, a.resnet3, a.smooth, a.rcut, + a.rcut_smth, a.protection, a.inv_nnei, a.edge_vec.data_ptr(), + a.edge_index.data_ptr(), a.edge_mask.data_ptr(), + a.atype.data_ptr(), a.davg.data_ptr(), + a.inv_dstd.data_ptr(), a.w1.data_ptr(), + a.w2.data_ptr(), optional_ptr(a.b2), optional_ptr(a.idt1), + optional_ptr(a.idt2), a.w3.data_ptr(), optional_ptr(a.b3), + optional_ptr(a.idt3), a.pair_table.data_ptr(), + optional_ptr(a.gate_table), a.order.data_ptr(), + dgr.data_ptr(), pre2_saved.data_ptr(), + g_saved.data_ptr(), d_edge_vec.data_ptr(), + pre2_saved.size(1)); + DPA1_CHECK_LAUNCH("dpa1_graph_descriptor backward"); +} + +// The backward uses four edges per thread once N1 reaches 16 to bound the +// register footprint of widening stacks. The N1 == 8 specialisation keeps +// eight edges per thread to avoid doubling the tile count. NG >= 128 +// prefetches saved g / pre3 rows into the available registers. +constexpr int backward_edges_per_thread(int n1) { return n1 >= 16 ? 4 : 8; } +constexpr int backward_prefetch(int ng) { return ng >= 128 ? 1 : 0; } + +} // namespace + +// Width / activation instantiation table shared by the two entry points: +// every "equal or doubling" stack over N1 in {8, 16, 32, 64} within the +// N2 <= 64, NG <= 128 tile bounds (see check_widths). Each stack expands to a +// runtime branch that resolves the activation and the tebd input mode (strip +// carries the type-pair gate, concat does not) to compile-time template +// arguments; the backward's per-thread edge fragment and spill prefetch follow +// from the width via backward_edges_per_thread / backward_prefetch. +#define DPA1_DISPATCH_ONE(LAUNCH, N1V, N2V, NGV) \ + else if (widths.n1 == N1V && widths.n2 == N2V && widths.ng == NGV) { \ + if (strip) { \ + if (act == 0) \ + LAUNCH(N1V, N2V, NGV, 0, 1); \ + else \ + LAUNCH(N1V, N2V, NGV, 1, 1); \ + } else { \ + if (act == 0) \ + LAUNCH(N1V, N2V, NGV, 0, 0); \ + else \ + LAUNCH(N1V, N2V, NGV, 1, 0); \ + } \ + } + +#define DPA1_DISPATCH_WIDTH_ACT(LAUNCH) \ + do { \ + if (false) { \ + } \ + DPA1_DISPATCH_ONE(LAUNCH, 8, 8, 8) \ + DPA1_DISPATCH_ONE(LAUNCH, 8, 8, 16) \ + DPA1_DISPATCH_ONE(LAUNCH, 8, 16, 16) \ + DPA1_DISPATCH_ONE(LAUNCH, 8, 16, 32) \ + DPA1_DISPATCH_ONE(LAUNCH, 16, 16, 16) \ + DPA1_DISPATCH_ONE(LAUNCH, 16, 16, 32) \ + DPA1_DISPATCH_ONE(LAUNCH, 16, 32, 32) \ + DPA1_DISPATCH_ONE(LAUNCH, 16, 32, 64) \ + DPA1_DISPATCH_ONE(LAUNCH, 32, 32, 32) \ + DPA1_DISPATCH_ONE(LAUNCH, 32, 32, 64) \ + DPA1_DISPATCH_ONE(LAUNCH, 32, 64, 64) \ + DPA1_DISPATCH_ONE(LAUNCH, 32, 64, 128) \ + DPA1_DISPATCH_ONE(LAUNCH, 64, 64, 64) \ + DPA1_DISPATCH_ONE(LAUNCH, 64, 64, 128) \ + else { \ + TORCH_CHECK(false, "dpa1_graph_descriptor: width dispatch miss"); \ + } \ + } while (0) + +// Forward: (grrg, rot_mat) plus the tensors the backward consumes. See the +// file header for the layout invariants and the applicability gate; the +// Python wrapper (deepmd.kernels.cuda.dpa1.graph_descriptor) documents the +// argument contract. An empty gate_table selects concat mode; a populated +// one ((T or T^2, NG), the strip embedding of the type pairs) selects strip. +std::tuple +dpa1_graph_descriptor(torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor atype, + torch::Tensor type_embedding, + torch::Tensor davg, + torch::Tensor dstd, + torch::Tensor w1, + torch::Tensor b1, + torch::Tensor idt1, + torch::Tensor w2, + torch::Tensor b2, + torch::Tensor idt2, + torch::Tensor w3, + torch::Tensor b3, + torch::Tensor idt3, + torch::Tensor gate_table, + int64_t act, + int64_t type_one_side, + int64_t concat_tebd, + int64_t write_rotation, + int64_t smooth, + int64_t axis, + int64_t resnet2, + int64_t resnet3, + double rcut, + double rcut_smth, + double protection, + double nnei) { + const auto widths = check_widths(w1, w2, w3); + const long n_edge = edge_vec.size(0); + const long n_node = atype.size(0); + const int ntypes = (int)type_embedding.size(0); + const int tebd_dim = (int)type_embedding.size(1); + const int NG = widths.ng; + const bool strip = gate_table.numel() > 0; + TORCH_CHECK(edge_vec.is_contiguous() && edge_index.is_contiguous() && + edge_mask.is_contiguous() && atype.is_contiguous(), + "dpa1_graph_descriptor: graph tensors must be contiguous"); + TORCH_CHECK( + (int)w1.size(0) == (strip ? 1 : 1 + (type_one_side ? 1 : 2) * tebd_dim), + "dpa1_graph_descriptor: w1 rows do not match the tebd mode"); + if (strip) { + TORCH_CHECK(gate_table.is_contiguous() && (int)gate_table.size(1) == NG && + gate_table.size(0) == + (type_one_side ? ntypes : (long)ntypes * ntypes), + "dpa1_graph_descriptor: gate_table must be a contiguous " + "(pairs, ng) strip embedding of the type pairs"); + } + TORCH_CHECK(edge_mask.scalar_type() == torch::kBool, + "dpa1_graph_descriptor: edge_mask must be bool"); + TORCH_CHECK(act == 0 || act == 1, + "dpa1_graph_descriptor: act must be 0 (tanh) or 1 (silu)"); + auto stream = at::cuda::getCurrentCUDAStream(); + + auto [order, pair_table] = build_order_and_pair_table( + edge_index, edge_mask, n_edge, n_node, ntypes, tebd_dim, widths.n1, + (int)type_one_side, (int)strip, type_embedding, w1, b1, stream); + auto inv_dstd = torch::reciprocal(dstd).contiguous(); + + auto f32 = + torch::TensorOptions().dtype(torch::kFloat32).device(edge_vec.device()); + auto gr = torch::zeros({n_node, 4, NG}, f32); + // Backward operands, spilled transposed so both directions stream + // coalesced rows. E_pad rounds up to the tile size so partial-tile float4 + // stores stay in bounds; for silu the g slot holds pre3 (its derivative + // needs the pre-activation). + const long e_pad = ceil_div(n_edge, kTileEdges) * kTileEdges; + auto pre2_saved = torch::empty({(long)widths.n2, e_pad}, f32); + auto g_saved = torch::empty({(long)NG, e_pad}, f32); + // The operator computes in fp32 (fp32 weights and tables); the coordinate + // input, in the model's global precision, enters through a single fp32 cast + // at the boundary (a no-op when already fp32). + const auto edge_vec_f = edge_vec.to(torch::kFloat32); + const LaunchArgs args{n_edge, + ntypes, + (int)type_one_side, + (int)resnet2, + (int)resnet3, + (int)smooth, + (float)rcut, + (float)rcut_smth, + (float)protection, + (float)(1.0 / nnei), + edge_vec_f, + edge_index, + edge_mask, + atype, + davg, + inv_dstd, + w1, + w2, + b2, + idt1, + idt2, + w3, + b3, + idt3, + pair_table, + gate_table, + order, + stream}; + +#define DPA1_LAUNCH_FWD(N1, N2, NG, ACT, STRIP) \ + launch_forward_portable(args, gr, pre2_saved, g_saved) + if (n_edge > 0) { + DPA1_DISPATCH_WIDTH_ACT(DPA1_LAUNCH_FWD); + } +#undef DPA1_LAUNCH_FWD + + const int out_dim = NG * (int)axis + (concat_tebd ? tebd_dim : 0); + auto grrg = torch::empty({n_node, out_dim}, f32); + auto rot_mat = torch::empty({write_rotation ? n_node : 0, NG, 3}, f32); + if (n_node > 0) { + gram_kernel<<<(int)n_node, 128, 4 * NG * sizeof(float), stream>>>( + (int)n_node, NG, (int)axis, tebd_dim, (int)concat_tebd, + gr.data_ptr(), type_embedding.data_ptr(), + atype.data_ptr(), grrg.data_ptr(), + write_rotation ? rot_mat.data_ptr() : nullptr); + DPA1_CHECK_LAUNCH("dpa1_graph_descriptor gram"); + } + return {grrg, rot_mat, gr, order, pair_table, pre2_saved, g_saved}; +} + +// Backward: (d_grrg, d_rot_mat) and the saved tensors -> d_edge_vec in the +// edge_vec dtype. rot_mat may be unused downstream (the energy fitting reads +// only grrg), in which case autograd passes None for its gradient. +torch::Tensor dpa1_graph_descriptor_backward( + torch::Tensor d_grrg, + std::optional d_rot_mat, + torch::Tensor gr, + torch::Tensor order, + torch::Tensor pair_table, + torch::Tensor pre2_saved, + torch::Tensor g_saved, + torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor atype, + torch::Tensor davg, + torch::Tensor dstd, + torch::Tensor w1, + torch::Tensor b1, + torch::Tensor idt1, + torch::Tensor w2, + torch::Tensor b2, + torch::Tensor idt2, + torch::Tensor w3, + torch::Tensor b3, + torch::Tensor idt3, + torch::Tensor gate_table, + int64_t act, + int64_t type_one_side, + int64_t smooth, + int64_t axis, + int64_t resnet2, + int64_t resnet3, + double rcut, + double rcut_smth, + double protection, + double nnei) { + const auto widths = check_widths(w1, w2, w3); + const long n_edge = edge_vec.size(0); + const long n_node = atype.size(0); + const int NG = widths.ng; + const bool strip = gate_table.numel() > 0; + auto stream = at::cuda::getCurrentCUDAStream(); + auto f32 = + torch::TensorOptions().dtype(torch::kFloat32).device(edge_vec.device()); + + auto dgr = torch::empty({n_node, 4, NG}, f32); + auto d_grrg_c = d_grrg.to(torch::kFloat32).contiguous(); + torch::Tensor d_rot_c; + const float* d_rot_ptr = nullptr; + if (d_rot_mat.has_value() && d_rot_mat->defined() && d_rot_mat->numel()) { + d_rot_c = d_rot_mat->to(torch::kFloat32).contiguous(); + d_rot_ptr = d_rot_c.data_ptr(); + } + if (n_node > 0) { + const size_t smem = (4 * NG + NG * (int)axis) * sizeof(float); + gram_backward_kernel<<<(int)n_node, 128, smem, stream>>>( + (int)n_node, NG, (int)axis, (int)d_grrg_c.size(1), + d_grrg_c.data_ptr(), d_rot_ptr, gr.data_ptr(), + dgr.data_ptr()); + DPA1_CHECK_LAUNCH("dpa1_graph_descriptor gram backward"); + } + + auto inv_dstd = torch::reciprocal(dstd).contiguous(); + // fp32 compute: cast the coordinate input in and produce the gradient in + // fp32; it is cast back to the model's precision at the boundary (see the + // forward). + const auto edge_vec_f = edge_vec.to(torch::kFloat32); + auto d_edge_vec = torch::empty({n_edge, 3}, f32); + // ntypes is recovered from the type-pair table row count (T or T^2): the + // layer-1 table in concat mode, the gate table in strip mode (whose + // layer-1 table has a single row). + const long n_pairs = strip ? gate_table.size(0) : pair_table.size(0); + const int ntypes = + type_one_side ? (int)n_pairs : (int)llround(sqrt((double)n_pairs)); + const LaunchArgs args{n_edge, + ntypes, + (int)type_one_side, + (int)resnet2, + (int)resnet3, + (int)smooth, + (float)rcut, + (float)rcut_smth, + (float)protection, + (float)(1.0 / nnei), + edge_vec_f, + edge_index, + edge_mask, + atype, + davg, + inv_dstd, + w1, + w2, + b2, + idt1, + idt2, + w3, + b3, + idt3, + pair_table, + gate_table, + order, + stream}; + +#define DPA1_LAUNCH_BWD(N1, N2, NG, ACT, STRIP) \ + launch_backward(args, dgr, pre2_saved, g_saved, \ + d_edge_vec) + if (n_edge > 0) { + DPA1_DISPATCH_WIDTH_ACT(DPA1_LAUNCH_BWD); + } +#undef DPA1_LAUNCH_BWD + // Return the gradient in the coordinate's precision (a no-op when fp32). + return d_edge_vec.to(edge_vec.scalar_type()); +} + +TORCH_LIBRARY_FRAGMENT(deepmd, m) { + m.def( + "dpa1_graph_descriptor(Tensor edge_vec, Tensor edge_index, " + "Tensor edge_mask, Tensor atype, Tensor type_embedding, Tensor davg, " + "Tensor dstd, Tensor w1, Tensor b1, Tensor idt1, Tensor w2, Tensor b2, " + "Tensor idt2, Tensor w3, Tensor b3, Tensor idt3, Tensor gate_table, " + "int act, int type_one_side, int concat_tebd, int write_rotation, int " + "smooth, int axis, int resnet2, int resnet3, float rcut, float " + "rcut_smth, " + "float protection, float nnei) -> (Tensor grrg, Tensor rot_mat, " + "Tensor gr, Tensor edge_order, Tensor pair_table, Tensor pre2_saved, " + "Tensor g_saved)"); + m.impl("dpa1_graph_descriptor", torch::kCUDA, &dpa1_graph_descriptor); + m.def( + "dpa1_graph_descriptor_backward(Tensor d_grrg, Tensor? d_rot_mat, " + "Tensor gr, Tensor edge_order, Tensor pair_table, Tensor pre2_saved, " + "Tensor g_saved, Tensor edge_vec, Tensor edge_index, Tensor edge_mask, " + "Tensor atype, Tensor davg, Tensor dstd, Tensor w1, Tensor b1, " + "Tensor idt1, Tensor w2, Tensor b2, Tensor idt2, Tensor w3, Tensor b3, " + "Tensor idt3, Tensor gate_table, int act, int type_one_side, " + "int smooth, int axis, int resnet2, int resnet3, float rcut, " + "float rcut_smth, float protection, float nnei) -> Tensor"); + m.impl("dpa1_graph_descriptor_backward", torch::kCUDA, + &dpa1_graph_descriptor_backward); +} diff --git a/source/op/pt/dpa1_graph_energy_force.cu b/source/op/pt/dpa1_graph_energy_force.cu new file mode 100644 index 0000000000..d1e9d56d3f --- /dev/null +++ b/source/op/pt/dpa1_graph_energy_force.cu @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// End-to-end fused energy-force operator for the DPA1 (``se_atten``) graph +// lower: one opaque graph node that consumes the neighbor edge stream and +// returns the per-frame energy, per-atom energy, force, virial and (optional) +// atom virial. It drives the descriptor and energy-fitting mega kernels and +// their analytic backwards in sequence, computing the force from the reduced +// energy internally (dE_redu/d(atom_energy) == 1), so the graph carries no +// autograd tape and no per-output-component grad loop. +// +// It is numerically identical to the separate-operator path (the descriptor +// and fitting forwards with the force assembled from autograd.grad): the same +// operator backwards evaluate the same fp32 arithmetic in the same order. The +// fusion removes the autograd machinery and the inter-operator array glue, and +// is selected at freeze time by DP_CUDA_INFER >= 2. + +#include +#include + +#include +#include +#include + +#include "graph_ops.h" + +// Returns (energy (nf, 1) fp64, atom_energy (N, 1) fp64, force (N, 3), +// virial (nf, 3, 3), atom_virial (N, 3, 3) or empty). +std::tuple +dpa1_graph_energy_force(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 atype, + torch::Tensor n_node, + torch::Tensor ownership, + torch::Tensor type_embedding, + torch::Tensor davg, + torch::Tensor dstd, + torch::Tensor w1, + torch::Tensor b1, + torch::Tensor idt1, + torch::Tensor w2, + torch::Tensor b2, + torch::Tensor idt2, + torch::Tensor w3, + torch::Tensor b3, + torch::Tensor idt3, + torch::Tensor gate_table, + int64_t act, + int64_t type_one_side, + int64_t concat_tebd, + int64_t smooth, + int64_t axis, + int64_t resnet2, + int64_t resnet3, + double rcut, + double rcut_smth, + double protection, + double nnei, + std::vector fit_ws, + std::vector fit_bs, + std::vector fit_idts, + std::vector fit_resnets, + torch::Tensor w_head, + torch::Tensor b_head, + torch::Tensor bias_atom_e, + int64_t fit_act, + c10::SymInt node_capacity, + bool do_atomic_virial) { + // The descriptor / fitting kernels compute in the weight precision (fp32). + // Cast the fp64 edge stream once and thread it through the forward, backward + // and force / virial scatter: each sub-operator's internal ``to(float)`` is + // then a no-op, and the backward returns the ``edge_vec`` gradient in that + // same precision, so the force scatter needs no cast either. Passing the + // fp64 leaf instead would re-cast the whole edge stream three times and + // round-trip the gradient fp32 -> fp64 -> fp32. The energy reduction stays + // fp64. + const auto fprec = w1.scalar_type(); + const auto edge_vec_f = + edge_vec.scalar_type() == fprec ? edge_vec : edge_vec.to(fprec); + + // === Step 1. Descriptor forward: edge stream -> (N, nd) descriptor. === + auto desc = dpa1_graph_descriptor( + edge_vec_f, edge_index, edge_mask, atype, type_embedding, davg, dstd, w1, + b1, idt1, w2, b2, idt2, w3, b3, idt3, gate_table, act, type_one_side, + concat_tebd, /*write_rotation=*/0, smooth, axis, resnet2, resnet3, rcut, + rcut_smth, protection, nnei); + const torch::Tensor& grrg = std::get<0>(desc); + const torch::Tensor& gr = std::get<2>(desc); + const torch::Tensor& edge_order = std::get<3>(desc); + const torch::Tensor& pair_table = std::get<4>(desc); + const torch::Tensor& pre2_saved = std::get<5>(desc); + 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); + 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()); + auto energy_seed = owned; + auto atom_energy = atom_energy_raw * owned; + + // === Step 3. Per-frame energy: segment-sum over the frame index. === + const int64_t nf = n_node.size(0); + auto frame_id = + at::repeat_interleave(at::arange(nf, n_node.options()), n_node); + auto energy = at::zeros({nf, 1}, atom_energy.options()) + .index_add_(0, frame_id, atom_energy); + + // === 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); + 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, + edge_vec_f, edge_index, edge_mask, atype, davg, dstd, w1, b1, idt1, w2, + b2, idt2, w3, b3, idt3, gate_table, act, type_one_side, smooth, axis, + resnet2, resnet3, rcut, rcut_smth, protection, nnei); + + // === 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); + return {energy, atom_energy, std::get<0>(fv), std::get<2>(fv), + std::get<1>(fv)}; +} + +TORCH_LIBRARY_FRAGMENT(deepmd, m) { + m.def( + "dpa1_graph_energy_force(Tensor edge_vec, Tensor edge_index, Tensor " + "edge_mask, Tensor destination_order, Tensor destination_row_ptr, " + "Tensor source_order, Tensor source_row_ptr, Tensor atype, Tensor " + "n_node, Tensor ownership, Tensor type_embedding, " + "Tensor davg, Tensor dstd, Tensor w1, Tensor b1, Tensor idt1, Tensor w2, " + "Tensor " + "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, Tensor[] fit_ws, Tensor[] fit_bs, Tensor[] fit_idts, 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/edge_force_virial.cu b/source/op/pt/edge_force_virial.cu new file mode 100644 index 0000000000..af0eb67861 --- /dev/null +++ b/source/op/pt/edge_force_virial.cu @@ -0,0 +1,478 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Force and virial assembly for an edge graph with two CSR views. +// +// Both CSR views store permutations into the original edge payload. One warp +// owns one node and reduces both incidence lists: +// +// force[node] = sum(dst=node) g_e - sum(src=node) g_e +// atom_virial[node] = sum(src=node) -g_e (x) edge_vec +// +// Every node therefore writes its force and atom virial exactly once; the hot +// path contains no global floating-point atomics. Per-frame virials are reduced +// from the per-node values through two FP64 stages and cast to the model +// precision only at the output boundary. + +#include +#include +#include + +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kWarpsPerBlock = kThreads / 32; +constexpr int kMaximumVirialPartials = 1024; + +#define FORCE_CHECK_LAUNCH(name) \ + do { \ + const cudaError_t error = cudaGetLastError(); \ + TORCH_CHECK(error == cudaSuccess, name, ": ", cudaGetErrorString(error)); \ + } while (0) + +__global__ void build_source_order_kernel(long valid_edge_count, + long edge_count, + const long* __restrict__ source, + long* __restrict__ cursor, + long* __restrict__ source_order) { + for (long edge = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + edge < edge_count; edge += static_cast(blockDim.x) * gridDim.x) { + if (edge < valid_edge_count) { + const long position = atomicAdd( + reinterpret_cast(cursor + source[edge]), 1ULL); + source_order[position] = edge; + } else { + source_order[edge] = edge; + } + } +} + +template +__global__ void edge_force_virial_kernel( + long node_count, + const scalar_t* __restrict__ edge_gradient, + const scalar_t* __restrict__ edge_vec, + const bool* __restrict__ edge_mask, + const index_t* __restrict__ destination_order, + const long* __restrict__ destination_row_ptr, + const index_t* __restrict__ source_order, + const long* __restrict__ source_row_ptr, + scalar_t* __restrict__ force, + scalar_t* __restrict__ node_virial) { + constexpr unsigned kWarpMask = 0xffffffffu; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const long node = static_cast(blockIdx.x) * kWarpsPerBlock + warp; + + scalar_t destination_x = 0; + scalar_t destination_y = 0; + scalar_t destination_z = 0; + scalar_t source_x = 0; + scalar_t source_y = 0; + scalar_t source_z = 0; + scalar_t virial[9] = {}; + + if (node < node_count) { + const long destination_begin = destination_row_ptr[node]; + const long destination_end = destination_row_ptr[node + 1]; + for (long position = destination_begin + lane; position < destination_end; + position += 32) { + const long edge = destination_order + ? static_cast(destination_order[position]) + : position; + if (edge_mask && !edge_mask[edge]) { + continue; + } + destination_x += edge_gradient[edge * 3 + 0]; + destination_y += edge_gradient[edge * 3 + 1]; + destination_z += edge_gradient[edge * 3 + 2]; + } + + const long source_begin = source_row_ptr[node]; + const long source_end = source_row_ptr[node + 1]; + for (long position = source_begin + lane; position < source_end; + position += 32) { + const long edge = static_cast(source_order[position]); + if (edge_mask && !edge_mask[edge]) { + continue; + } + const scalar_t gx = edge_gradient[edge * 3 + 0]; + const scalar_t gy = edge_gradient[edge * 3 + 1]; + const scalar_t gz = edge_gradient[edge * 3 + 2]; + const scalar_t x = edge_vec[edge * 3 + 0]; + const scalar_t y = edge_vec[edge * 3 + 1]; + const scalar_t z = edge_vec[edge * 3 + 2]; + source_x += gx; + source_y += gy; + source_z += gz; + virial[0] = fma(-gx, x, virial[0]); + virial[1] = fma(-gx, y, virial[1]); + virial[2] = fma(-gx, z, virial[2]); + virial[3] = fma(-gy, x, virial[3]); + virial[4] = fma(-gy, y, virial[4]); + virial[5] = fma(-gy, z, virial[5]); + virial[6] = fma(-gz, x, virial[6]); + virial[7] = fma(-gz, y, virial[7]); + virial[8] = fma(-gz, z, virial[8]); + } + } + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + destination_x += __shfl_down_sync(kWarpMask, destination_x, offset); + destination_y += __shfl_down_sync(kWarpMask, destination_y, offset); + destination_z += __shfl_down_sync(kWarpMask, destination_z, offset); + 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); +#pragma unroll + for (int component = 0; component < 9; ++component) { + virial[component] += + __shfl_down_sync(kWarpMask, virial[component], offset); + } + } + + if (lane == 0 && node < node_count) { + force[node * 3 + 0] = destination_x - source_x; + force[node * 3 + 1] = destination_y - source_y; + force[node * 3 + 2] = destination_z - source_z; + scalar_t* output = node_virial + node * 9; +#pragma unroll + for (int component = 0; component < 9; ++component) { + output[component] = virial[component]; + } + } +} + +template +__global__ void reduce_node_virial_kernel( + long frame_count, + int partial_count, + const long* __restrict__ frame_row_ptr, + const scalar_t* __restrict__ node_virial, + double* __restrict__ partial) { + __shared__ double values[kThreads]; + const long task_count = static_cast(frame_count) * 9 * 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 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]); + } + values[threadIdx.x] = sum; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + values[threadIdx.x] += values[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + partial[output * partial_count + partial_index] = values[0]; + } + __syncthreads(); + } +} + +template +__global__ void finalize_virial_kernel(long output_count, + int partial_count, + const double* __restrict__ partial, + scalar_t* __restrict__ virial) { + __shared__ double values[kThreads]; + for (long output = blockIdx.x; output < output_count; output += gridDim.x) { + double sum = 0.0; + for (int index = threadIdx.x; index < partial_count; index += blockDim.x) { + sum += partial[output * partial_count + index]; + } + values[threadIdx.x] = sum; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + values[threadIdx.x] += values[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + virial[output] = static_cast(values[0]); + } + __syncthreads(); + } +} + +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, + 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& frame_row_ptr, + torch::Tensor& force, + torch::Tensor& node_virial, + 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()); + 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"); +} + +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) { + const long frame_count = n_node_per_frame.size(0); + auto options = edge_gradient.options(); + auto force = torch::empty({node_count, 3}, options); + auto atom_virial = + torch::empty({want_atom_virial ? node_count : 0, 3, 3}, options); + auto node_virial = want_atom_virial + ? atom_virial + : torch::empty({node_count, 3, 3}, options); + auto virial = torch::zeros({frame_count, 3, 3}, options); + if (node_count == 0 || frame_count == 0) { + return {force, atom_virial, virial}; + } + + auto frame_row_ptr = + torch::cat({torch::zeros({1}, n_node_per_frame.options()), + torch::cumsum(n_node_per_frame, 0)}) + .to(torch::kInt64) + .contiguous(); + const long average_node_count = (node_count + frame_count - 1) / frame_count; + const int partial_count = + static_cast(std::min((average_node_count + kThreads - 1) / kThreads, + static_cast(kMaximumVirialPartials))); + auto virial_partial = torch::empty({frame_count * 9, partial_count}, + options.dtype(torch::kFloat64)); + const auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES( + 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); + } 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); + } + }); + return {force, atom_virial, virial}; +} + +} // namespace + +std::tuple +build_graph_csr(torch::Tensor edge_index, + c10::SymInt node_count_symbol, + c10::SymInt valid_edge_count_symbol) { + const long node_count = node_count_symbol.expect_int(); + const long valid_edge_count = valid_edge_count_symbol.expect_int(); + const long edge_count = edge_index.size(1); + TORCH_CHECK(edge_index.is_cuda() && edge_index.is_contiguous(), + "build_graph_csr: edge_index must be a contiguous CUDA tensor"); + TORCH_CHECK(edge_index.scalar_type() == torch::kInt64, + "build_graph_csr: edge_index must be int64"); + TORCH_CHECK(node_count > 0, "build_graph_csr: node_count must be positive"); + TORCH_CHECK(valid_edge_count >= 0 && valid_edge_count <= edge_count, + "build_graph_csr: valid_edge_count must lie in [0, E]"); + + const auto valid_source = + edge_index.select(0, 0).slice(0, 0, valid_edge_count); + const auto valid_destination = + edge_index.select(0, 1).slice(0, 0, valid_edge_count); + const auto source_counts = torch::bincount(valid_source, {}, node_count); + const auto destination_counts = + torch::bincount(valid_destination, {}, node_count); + const auto zero = torch::zeros({1}, source_counts.options()); + auto source_row_ptr = torch::cat({zero, torch::cumsum(source_counts, 0)}) + .to(torch::kInt64) + .contiguous(); + auto destination_row_ptr = + torch::cat({zero, torch::cumsum(destination_counts, 0)}) + .to(torch::kInt64) + .contiguous(); + auto destination_order = torch::arange(edge_count, edge_index.options()); + auto source_order = torch::empty({edge_count}, edge_index.options()); + auto cursor = source_row_ptr.slice(0, 0, node_count).clone(); + + const int blocks = + std::min(static_cast((edge_count + kThreads - 1) / kThreads), 65535); + if (blocks > 0) { + const auto stream = at::cuda::getCurrentCUDAStream(); + build_source_order_kernel<<>>( + valid_edge_count, edge_count, edge_index.select(0, 0).data_ptr(), + cursor.data_ptr(), source_order.data_ptr()); + FORCE_CHECK_LAUNCH("build_graph_csr source order"); + } + 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) { + 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() && + destination_row_ptr.is_cuda() && source_order.is_cuda() && + source_row_ptr.is_cuda() && n_node_per_frame.is_cuda(), + "edge_force_virial: edge and CSR tensors must be CUDA tensors"); + TORCH_CHECK(edge_gradient.is_contiguous() && edge_vec.is_contiguous() && + edge_mask.is_contiguous(), + "edge_force_virial: edge tensors must be contiguous"); + TORCH_CHECK(edge_gradient.scalar_type() == edge_vec.scalar_type(), + "edge_force_virial: gradient and edge vector dtypes must match"); + TORCH_CHECK(edge_mask.scalar_type() == torch::kBool, + "edge_force_virial: edge_mask must be bool"); + TORCH_CHECK(destination_row_ptr.scalar_type() == torch::kInt64 && + source_row_ptr.scalar_type() == torch::kInt64, + "edge_force_virial: CSR row pointers must be int64"); + TORCH_CHECK(destination_order.is_contiguous() && + destination_row_ptr.is_contiguous() && + source_order.is_contiguous() && + source_row_ptr.is_contiguous(), + "edge_force_virial: CSR tensors must be contiguous"); + TORCH_CHECK( + (source_order.scalar_type() == torch::kInt32 || + 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"); + 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); +} + +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, + 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() && + destination_row_ptr.is_cuda() && source_row_ptr.is_cuda() && + source_order.is_cuda() && n_node_per_frame.is_cuda(), + "canonical_edge_force_virial: inputs must be CUDA tensors"); + TORCH_CHECK(edge_gradient.is_contiguous() && edge_vec.is_contiguous() && + destination_row_ptr.is_contiguous() && + source_row_ptr.is_contiguous() && + source_order.is_contiguous(), + "canonical_edge_force_virial: inputs must be contiguous"); + TORCH_CHECK(edge_gradient.sizes() == edge_vec.sizes() && + edge_gradient.scalar_type() == edge_vec.scalar_type(), + "canonical_edge_force_virial: gradient and edge vector must " + "share shape and dtype"); + TORCH_CHECK(destination_row_ptr.scalar_type() == torch::kInt64 && + 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::kInt64, + "canonical_edge_force_virial: source_order must be int32 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"); + + 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); +} + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.def( + "build_graph_csr(Tensor edge_index, SymInt node_count, " + "SymInt valid_edge_count) -> " + "(Tensor destination_order, Tensor destination_row_ptr, " + "Tensor source_order, Tensor source_row_ptr)"); + library.impl("build_graph_csr", torch::kCUDA, &build_graph_csr); + library.def( + "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)"); + 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, " + "bool want_atom_virial) -> " + "(Tensor force, Tensor atom_virial, Tensor virial)"); + library.impl("canonical_edge_force_virial", torch::kCUDA, + &canonical_edge_force_virial); +} + +TORCH_LIBRARY_IMPL(deepmd, Autograd, library) { + library.impl("edge_force_virial", torch::CppFunction::makeFallthrough()); + library.impl("canonical_edge_force_virial", + torch::CppFunction::makeFallthrough()); +} diff --git a/source/op/pt/graph_fitting.cu b/source/op/pt/graph_fitting.cu new file mode 100644 index 0000000000..60724679c8 --- /dev/null +++ b/source/op/pt/graph_fitting.cu @@ -0,0 +1,380 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// 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) +// 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 +// dh_{L-1} = d_e * w_head^T +// dpre_l = dh_l * act'_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. +// +// 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 + +namespace { + +#define FITTING_CHECK_LAUNCH(what) \ + do { \ + cudaError_t err = cudaGetLastError(); \ + TORCH_CHECK(err == cudaSuccess, what, ": ", cudaGetErrorString(err)); \ + } while (0) + +cublasHandle_t cublas_handle() { + // A cuBLAS handle is device-bound and unsafe to share across threads, so + // cache one per device in thread-local storage. Pedantic math keeps the fp32 + // potential-energy surface exact (no TF32, no split-K reordering) for MD. + thread_local std::unordered_map handles; + int device = 0; + cudaGetDevice(&device); + cublasHandle_t& h = handles[device]; + if (!h) { + cublasCreate(&h); + cublasSetMathMode(h, CUBLAS_PEDANTIC_MATH); + } + return h; +} + +// Row-major C(m, n) = A(m, k) @ B(k, n) + beta * C. +void gemm_nn(cudaStream_t stream, + const float* a, + const float* b, + float* c, + int m, + int n, + int k, + float beta = 0.f) { + cublasSetStream(cublas_handle(), stream); + const float alpha = 1.f; + cublasSgemm(cublas_handle(), CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, b, n, + a, k, &beta, c, n); +} + +// Row-major C(m, n) = A(m, k) @ B(n, k)^T + beta * C. +void gemm_nt(cudaStream_t stream, + const float* a, + const float* b, + float* c, + int m, + int n, + int k, + float beta = 0.f) { + cublasSetStream(cublas_handle(), stream); + const float alpha = 1.f; + cublasSgemm(cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, n, m, k, &alpha, b, k, + a, k, &beta, c, n); +} + +// Activation value and derivative; codes follow +// deepmd.kernels.triton.dpa1.activation.ACT_CODES (0 = tanh, 1 = silu). +template +__device__ __forceinline__ float2 act_vg(float z) { + if constexpr (ACT == 0) { + float a = tanhf(z); + return make_float2(a, 1.f - a * a); + } + const float s = 0.5f * (1.f + tanhf(0.5f * z)); // sigmoid via tanh identity + return make_float2(z * s, s * (1.f + z * (1.f - s))); +} + +// 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, + 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) { + return; + } + const long t = t4 * 4; + const int c = (int)(t % dout); + 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); + if (residual) { + const float4 xx = *reinterpret_cast(x + t); + yy.x += xx.x; + yy.y += xx.y; + yy.z += xx.z; + 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]]. +__global__ void head_kernel(long n_node, + int width, + const float* __restrict__ h, + const float* __restrict__ w_head, + const float* __restrict__ b_head, + const double* __restrict__ bias_atom_e, + const long* __restrict__ atype, + double* __restrict__ e) { + const long n = blockIdx.x * (long)blockDim.x + threadIdx.x; + if (n >= n_node) { + return; + } + const float* row = h + n * width; + float acc = 0.f; + for (int k = 0; k < width; k += 4) { + const float4 hv = *reinterpret_cast(row + k); + const float4 wv = *reinterpret_cast(w_head + k); + acc += hv.x * wv.x + hv.y * wv.y + hv.z * wv.z + hv.w * wv.w; + } + 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) { + 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 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); +} + +// Convert dh to dpre in place: dh *= adot. +__global__ void backward_epilogue_kernel(long total4, + const float* __restrict__ dh, + const float* __restrict__ adot, + float* __restrict__ dpre) { + const long t4 = blockIdx.x * (long)blockDim.x + threadIdx.x; + if (t4 >= total4) { + return; + } + const long t = t4 * 4; + const float4 d = *reinterpret_cast(dh + t); + const float4 a = *reinterpret_cast(adot + t); + *reinterpret_cast(dpre + t) = + make_float4(d.x * a.x, d.y * a.y, d.z * a.z, d.w * a.w); +} + +long ceil_div(long a, long b) { return (a + b - 1) / b; } + +} // namespace + +// 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. +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 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 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"); + return {e, saved}; +} + +// Backward: d_x from the upstream d_e (fp64 (N, 1)). The saved derivative +// 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. +void graph_fitting_backward_core(torch::Tensor d_e, + torch::Tensor saved, + std::vector ws, + std::vector resnets, + torch::Tensor w_head, + 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 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 && + d_x.scalar_type() == torch::kFloat32 && d_x.is_cuda() && + d_x.is_contiguous(), + "graph_fitting_backward: output must be contiguous CUDA " + "fp32 with shape (N, input_width)"); + // Guard the empty system before the division by ``n_node`` below. + 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); + } + 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); + } + } +} + +torch::Tensor graph_fitting_backward(torch::Tensor d_e, + torch::Tensor saved, + std::vector ws, + 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; + 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); + return d_x; +} + +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)"); + 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"); + m.impl("graph_fitting_backward", torch::kCUDA, &graph_fitting_backward); +} diff --git a/source/op/pt/graph_ops.h b/source/op/pt/graph_ops.h new file mode 100644 index 0000000000..669513b2b5 --- /dev/null +++ b/source/op/pt/graph_ops.h @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Entry points of the fused graph-lower operator suite, shared so the fused +// energy-force operator (dpa1_graph_energy_force.cu) can drive the forward and +// backward passes directly rather than through the dispatcher. The definitions +// live in dpa1_graph_descriptor.cu, graph_fitting.cu and edge_force_virial.cu; +// the operator schemas registered from those files are the public interface. +// The compressed descriptor (dpa1_graph_compress.cu) is not part of the +// end-to-end operator, so its entries stay private to that translation unit. + +#pragma once + +#include + +#include +#include +#include + +// DPA1 descriptor body (environment matrix, embedding MLP, moment, G^T G). +// Returns (grrg, rot_mat, gr, edge_order, pair_table, pre2_saved, g_saved); +// the last five are consumed by dpa1_graph_descriptor_backward. +std::tuple +dpa1_graph_descriptor(torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor atype, + torch::Tensor type_embedding, + torch::Tensor davg, + torch::Tensor dstd, + torch::Tensor w1, + torch::Tensor b1, + torch::Tensor idt1, + torch::Tensor w2, + torch::Tensor b2, + torch::Tensor idt2, + torch::Tensor w3, + torch::Tensor b3, + torch::Tensor idt3, + torch::Tensor gate_table, + int64_t act, + int64_t type_one_side, + int64_t concat_tebd, + int64_t write_rotation, + int64_t smooth, + int64_t axis, + int64_t resnet2, + int64_t resnet3, + double rcut, + double rcut_smth, + double protection, + double nnei); + +// dE/d(edge_vec) from dE/d(grrg); consumes the saved tensors of the forward. +torch::Tensor dpa1_graph_descriptor_backward( + torch::Tensor d_grrg, + std::optional d_rot_mat, + torch::Tensor gr, + torch::Tensor order, + torch::Tensor pair_table, + torch::Tensor pre2_saved, + torch::Tensor g_saved, + torch::Tensor edge_vec, + torch::Tensor edge_index, + torch::Tensor edge_mask, + torch::Tensor atype, + torch::Tensor davg, + torch::Tensor dstd, + torch::Tensor w1, + torch::Tensor b1, + torch::Tensor idt1, + torch::Tensor w2, + torch::Tensor b2, + torch::Tensor idt2, + torch::Tensor w3, + torch::Tensor b3, + torch::Tensor idt3, + torch::Tensor gate_table, + int64_t act, + int64_t type_one_side, + int64_t smooth, + int64_t axis, + int64_t resnet2, + int64_t resnet3, + double rcut, + double rcut_smth, + double protection, + double nnei); + +// Energy fitting MLP on the flat node axis. Returns (atom_energy (N, 1) fp64, +// saved activations/derivatives) 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. +torch::Tensor graph_fitting_backward(torch::Tensor d_e, + torch::Tensor saved, + std::vector ws, + std::vector resnets, + torch::Tensor w_head); + +// 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); diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion.py b/source/tests/common/dpmodel/test_apply_pair_exclusion.py index b73f4c8d38..9ff1dfb29b 100644 --- a/source/tests/common/dpmodel/test_apply_pair_exclusion.py +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion.py @@ -8,6 +8,7 @@ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, apply_pair_exclusion, + attach_edge_csr, ) @@ -46,6 +47,19 @@ def test_excluded_pairs_are_masked_and_padding_stays_masked() -> None: assert out.edge_vec is g.edge_vec +def test_exclusion_invalidates_derived_csr_views() -> None: + graph = attach_edge_csr(_toy_graph(), 4) + atype = np.array([0, 1, 0, 1], dtype=np.int64) + + out = apply_pair_exclusion(graph, atype, PairExcludeMask(2, [(0, 1)])) + + assert out.destination_order is None + assert out.destination_row_ptr is None + assert out.source_order is None + assert out.source_row_ptr is None + assert out.destination_sorted is False + + def test_no_exclusion_empty_list_is_identity() -> None: """Cover PairExcludeMask with non-None but empty exclude list.""" g = _toy_graph() diff --git a/source/tests/common/dpmodel/test_call_lower_graph.py b/source/tests/common/dpmodel/test_call_lower_graph.py index 29dd87c55a..dbedf4ee5b 100644 --- a/source/tests/common/dpmodel/test_call_lower_graph.py +++ b/source/tests/common/dpmodel/test_call_lower_graph.py @@ -3,9 +3,7 @@ (``CM.call_lower_graph``) and the dense ``EnergyModel.call_lower`` on the SAME neighbor list (regime-1: ``from_dense_quartet`` reproduces the nlist neighbors). -PR-A is dpa1(attn_layer=0) energy-only; force/virial come from pt_expt autograd -in a later task, so this only checks ``energy`` (reduced per-frame) and -``atom_energy`` (per-atom). +This suite checks ``energy`` (reduced per-frame) and ``atom_energy`` (per-atom). """ import unittest diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py index b889644abc..0f072c70b3 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py @@ -162,8 +162,8 @@ def test_exclude_types_graph_eligible_and_parity(self, exclude_types) -> None: def test_eligible_no_mapping_with_ghosts_falls_back(self) -> None: """An eligible (concat) attn_layer=0 descriptor called with mapping=None on a PERIODIC system (nall > nloc ghosts) must fall back to the dense - body and match it (regression: the graph needs mapping for ghosts, the - identity-mapping default previously indexed out of range). + body and match it because the graph path requires an explicit ghost + mapping. """ dd = self._make([30]) box = np.eye(3, dtype=np.float64)[None] * 6.0 diff --git a/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py b/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py index 4eec03ba0b..b156e0ff5b 100644 --- a/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py +++ b/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Graph-native se_atten attention (attn_layer > 0) vs the dense reference. -Regime-1 parity (NeighborGraph PR-D): the graph is built FROM the same dense +The graph is built from the same dense nlist (``from_dense_quartet``), so the neighbor sets are identical and the graph attention must reproduce ``GatedAttentionLayer``/``NeighborGatedAttention`` -bit-exactly (CPU rtol 1e-12) for ANY sel — binding or not. +bit-exactly (CPU rtol 1e-12) for any sel. The smooth branch needs the SHAPE-STATIC graph (``compact=False`` + ``static_nnei``): dense smooth keeps padding slots in the softmax DENOMINATOR @@ -137,7 +137,6 @@ def test_block_compact_graph_no_smooth(self, attn_layer) -> None: graph_g.reshape(dense_g.shape), dense_g, rtol=1e-12, atol=1e-12 ) - # ── smooth on the compact (carry-all) form: CLEAN DIVERGENCE by design ──── def test_block_compact_graph_smooth_clean_divergence(self) -> None: """Carry-all smooth attention deliberately DIVERGES from dense. @@ -145,8 +144,7 @@ def test_block_compact_graph_smooth_clean_divergence(self) -> None: softmax DENOMINATOR at weight ``exp(-attnw_shift)``, which makes the dense output depend on ``sel`` itself (same physical neighbors, different sel => different output, up to ~1e-4). The carry-all graph - drops those phantom terms — the sel-independent math (user decision - 2026-07-03, PR-D). Bit-parity (1e-12) is proven on the shape-static + drops those phantom terms. Bit-parity (1e-12) is proven on the shape-static adapter (same padded pairs on both sides, ``test_dotr_smooth``); here we pin only that the compact form stays CLOSE to dense (the artifact is a bounded denominator perturbation) while NOT bit-equal. @@ -178,8 +176,7 @@ def test_block_compact_graph_smooth_clean_divergence(self) -> None: # ... but NOT bit-equal: the phantom-padding terms are really gone assert np.max(np.abs(graph_g - dense_g)) > 1e-9 - # ── torch namespace smoke (CLAUDE.md: catches numpy-weight leaks) ──────── - # NB: the smoke runs the BLOCK kernel with a torch type_embedding table; + # The smoke test runs the block kernel with a torch type_embedding table; # the raw dpmodel adapter is numpy-weighted by design (pt_expt wraps it). def test_torch_block_matches_numpy(self) -> None: import torch diff --git a/source/tests/common/dpmodel/test_dpa1_graph_model_energy.py b/source/tests/common/dpmodel/test_dpa1_graph_model_energy.py index 37cb18808e..e03dcaa840 100644 --- a/source/tests/common/dpmodel/test_dpa1_graph_model_energy.py +++ b/source/tests/common/dpmodel/test_dpa1_graph_model_energy.py @@ -227,8 +227,8 @@ def test_neighbor_list_takes_dense_route() -> None: def test_graph_lower_invariant_to_charge_spin() -> None: """dpa1 does NOT consume charge_spin (``get_dim_chg_spin() == 0``); the dense atomic model passes ``None`` to the dpa1 descriptor regardless. The graph - lower accepts ``charge_spin`` only for ABI stability with charge/spin - descriptors (dpa3/dpa4, PR-G), so its output must be INVARIANT to it. + lower accepts ``charge_spin`` for descriptors with charge/spin + conditioning, so dpa1 output must be invariant to it. Combined with the graph==dense parity at non-binding sel (:func:`test_energy_parity_non_binding_sel`), this gives the full claim: diff --git a/source/tests/common/dpmodel/test_edge_force_virial.py b/source/tests/common/dpmodel/test_edge_force_virial.py index 722960f57a..7b773ab1d2 100644 --- a/source/tests/common/dpmodel/test_edge_force_virial.py +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -97,19 +97,30 @@ def test_all_edges_masked_gives_zero(self) -> None: np.testing.assert_allclose(av, np.zeros((n, 3, 3))) np.testing.assert_allclose(vir, np.zeros((nf, 3, 3))) + def test_zero_nodes_with_masked_guards_gives_empty_outputs(self) -> None: + edge_index = np.zeros((2, 2), dtype=np.int64) + edge_vec = np.zeros((2, 3)) + edge_mask = np.zeros(2, dtype=np.bool_) + g = np.ones((2, 3)) + n_node = np.array([0], dtype=np.int64) + + force, atom_virial, virial = edge_force_virial( + g, + edge_vec, + edge_index, + edge_mask, + n_node, + ) + + self.assertEqual(force.shape, (0, 3)) + self.assertEqual(atom_virial.shape, (0, 3, 3)) + self.assertEqual(virial.shape, (1, 3, 3)) + np.testing.assert_allclose(virial, 0.0) + def test_modulo_clamp_leaves_real_edges_unchanged(self) -> None: - # INVARIANT (iProzd review): the in-bounds index clamp (``% n_out``) that - # keeps the CUDA-exported scatter address legal must NEVER alter a real - # (edge_mask == True) edge -- only masked/out-of-range guard edges may be - # remapped, and those carry zero weight so remapping is harmless. Here a - # REAL edge sits on the boundary node ``n_out - 1`` (the largest valid - # index, where a wrong wrap would be visible) and a MASKED guard edge - # carries deliberately OUT-OF-RANGE indices (>= n_out) with nonzero g/vec. - # Correctness requires the result to equal the real-edges-only reference: - # the boundary real edge must land on node n_out-1 (not wrapped), and the - # out-of-range guard must contribute nothing. If real edges were ever - # remapped by the modulo (the shape-binding bug iProzd warned about), the - # boundary node's force/virial would be wrong and this test would fail. + # The endpoint bound must leave every valid edge unchanged. Masked + # guards may carry out-of-range sentinels because their contributions + # are zeroed before the scatter. n_node = np.array([5], dtype=np.int64) # 1 frame, nodes 0..4 (n_out = 5) # e0: real, src on the boundary node 4 -> node 0 ; e1: real, node 0 -> 4 # e2: MASKED guard with out-of-range indices src=99, dst=77 (>= n_out) diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index ef8066850b..b8b57c8b31 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -6,6 +6,8 @@ from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, NeighborGraph, + build_edge_csr, + canonicalize_neighbor_graph, ) @@ -20,6 +22,7 @@ def test_construct_minimal(self) -> None: self.assertEqual(ng.edge_index.shape, (2, 2)) self.assertEqual(ng.edge_vec.shape, (2, 3)) # optionals default to None + self.assertFalse(ng.destination_sorted) self.assertIsNone(ng.n_local) self.assertIsNone(ng.angle_index) self.assertIsNone(ng.angle_mask) @@ -33,7 +36,114 @@ def test_graphlayout_defaults(self) -> None: self.assertEqual(lay.min_edges, 2) +class TestBuildEdgeCSR(unittest.TestCase): + def test_preserves_payload_and_builds_stable_permutations(self) -> None: + edge_vec = np.arange(15, dtype=np.float64).reshape(5, 3) + edge_mask = np.array([True, False, True, True, False]) + for index_dtype in (np.int32, np.int64): + with self.subTest(index_dtype=index_dtype): + edge_index = np.array( + [[2, 0, 1, 2, 0], [1, 2, 0, 0, 1]], + dtype=index_dtype, + ) + ( + result_index, + result_vec, + result_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr(edge_index, edge_vec, edge_mask, n_nodes=3) + + np.testing.assert_array_equal(result_index, edge_index) + np.testing.assert_array_equal(result_vec, edge_vec) + np.testing.assert_array_equal(result_mask, edge_mask) + np.testing.assert_array_equal( + destination_order, np.array([2, 3, 0, 1, 4]) + ) + np.testing.assert_array_equal( + destination_row_ptr, np.array([0, 2, 3, 3]) + ) + np.testing.assert_array_equal(source_order, np.array([2, 0, 3, 1, 4])) + np.testing.assert_array_equal(source_row_ptr, np.array([0, 0, 1, 3])) + self.assertEqual(destination_order.dtype, index_dtype) + self.assertEqual(source_order.dtype, index_dtype) + self.assertEqual(destination_row_ptr.dtype, np.int64) + self.assertEqual(source_row_ptr.dtype, np.int64) + + def test_canonicalizes_destination_payload_stably(self) -> None: + edge_index = np.array([[2, 0, 1, 2, 0], [1, 2, 0, 0, 1]], dtype=np.int64) + edge_vec = np.arange(15, dtype=np.float64).reshape(5, 3) + edge_mask = np.array([True, False, True, True, False]) + ( + result_index, + result_vec, + result_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + edge_index, + edge_vec, + edge_mask, + n_nodes=3, + canonicalize=True, + ) + + expected_permutation = np.array([2, 3, 0, 1, 4]) + np.testing.assert_array_equal(result_index, edge_index[:, expected_permutation]) + np.testing.assert_array_equal(result_vec, edge_vec[expected_permutation]) + np.testing.assert_array_equal( + result_mask, np.array([True, True, True, False, False]) + ) + np.testing.assert_array_equal(destination_order, np.arange(5)) + np.testing.assert_array_equal(destination_row_ptr, np.array([0, 2, 3, 3])) + np.testing.assert_array_equal(source_order, np.arange(5)) + np.testing.assert_array_equal(source_row_ptr, np.array([0, 0, 1, 3])) + + def test_canonicalizes_neighbor_graph_at_deployment_boundary(self) -> None: + edge_index = np.array([[2, 1, 0], [1, 0, 2]], dtype=np.int64) + graph = NeighborGraph( + n_node=np.array([3], dtype=np.int64), + edge_index=edge_index, + edge_vec=np.arange(9, dtype=np.float64).reshape(3, 3), + edge_mask=np.array([True, True, True]), + ) + + result = canonicalize_neighbor_graph(graph, n_nodes=3) + + self.assertTrue(result.destination_sorted) + np.testing.assert_array_equal( + result.edge_index, edge_index[:, np.array([1, 0, 2])] + ) + np.testing.assert_array_equal(result.destination_order, np.arange(3)) + + def test_canonicalization_rebuilds_an_untrusted_canonical_claim(self) -> None: + graph = NeighborGraph( + n_node=np.array([2], dtype=np.int64), + edge_index=np.array([[0, 1], [1, 0]], dtype=np.int64), + edge_vec=np.zeros((2, 3), dtype=np.float64), + edge_mask=np.array([True, True]), + destination_order=np.array([1, 0], dtype=np.int64), + destination_row_ptr=np.array([0, 1, 2], dtype=np.int64), + source_order=np.array([0, 1], dtype=np.int64), + source_row_ptr=np.array([0, 1, 2], dtype=np.int64), + destination_sorted=True, + ) + + result = canonicalize_neighbor_graph(graph, n_nodes=2) + + np.testing.assert_array_equal(result.destination_order, np.arange(2)) + np.testing.assert_array_equal( + result.edge_index, + np.array([[1, 0], [0, 1]], dtype=np.int64), + ) + + from deepmd.dpmodel.utils.neighbor_graph import ( + node_ownership_mask, node_validity_mask, ) @@ -49,6 +159,15 @@ def test_with_padding_prefix(self) -> None: mask = node_validity_mask(n_node, 8) # N_max = 8 => 3 padding np.testing.assert_array_equal(mask, np.array([True] * 5 + [False] * 3)) + def test_local_plus_halo_ownership(self) -> None: + n_node = np.array([3, 4], dtype=np.int64) + n_local = np.array([2, 1], dtype=np.int64) + mask = node_ownership_mask(n_node, n_local, 7) + np.testing.assert_array_equal( + mask, + np.array([True, True, False, True, False, False, False]), + ) + from deepmd.dpmodel.utils.neighbor_graph import ( pad_and_guard_edges, diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index 8f91f5af8c..4f45d22505 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -104,6 +104,45 @@ def test_nonperiodic_matches_brute_force(self) -> None: brute_force_neighbor_sets(self.coord[0], None, self.rcut), ) + def test_canonicalization_is_explicit(self) -> None: + generic = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + with_csr = build_neighbor_graph( + self.coord, + self.atype, + None, + self.rcut, + with_csr=True, + ) + canonical = build_neighbor_graph( + self.coord, + self.atype, + None, + self.rcut, + canonicalize=True, + ) + + self.assertFalse(generic.destination_sorted) + self.assertIsNone(generic.destination_order) + self.assertIsNone(generic.destination_row_ptr) + self.assertIsNone(generic.source_order) + self.assertIsNone(generic.source_row_ptr) + self.assertFalse(with_csr.destination_sorted) + self.assertIsNotNone(with_csr.destination_order) + self.assertIsNotNone(with_csr.destination_row_ptr) + self.assertIsNotNone(with_csr.source_order) + self.assertIsNotNone(with_csr.source_row_ptr) + self.assertTrue(canonical.destination_sorted) + np.testing.assert_array_equal( + canonical.destination_order, + np.arange(canonical.edge_index.shape[1]), + ) + real_destination = canonical.edge_index[1, canonical.edge_mask] + self.assertTrue(bool(np.all(real_destination[:-1] <= real_destination[1:]))) + self.assertEqual( + graph_neighbor_sets(canonical, 4), + graph_neighbor_sets(generic, 4), + ) + def test_periodic_matches_brute_force(self) -> None: box = np.eye(3, dtype=np.float64)[None] * 6.0 ng = build_neighbor_graph(self.coord, self.atype, box, self.rcut) @@ -382,6 +421,25 @@ def test_oracle_set_equality_dense_nonperiodic(self) -> None: # sanity: exclusion actually REMOVED some edges self.assertLess(int(ng_fused.edge_mask.sum()), int(ng_base.edge_mask.sum())) + def test_canonical_csr_reflects_pair_exclusion(self) -> None: + pe = self._pair_excl([(0, 1), (1, 0)]) + graph = build_neighbor_graph( + self.coord, + self.atype, + None, + self.rcut, + pair_excl=pe, + canonicalize=True, + ) + + valid_edges = int(graph.edge_mask.sum()) + self.assertEqual(int(graph.destination_row_ptr[-1]), valid_edges) + self.assertEqual(int(graph.source_row_ptr[-1]), valid_edges) + self.assertTrue(bool(np.all(graph.edge_mask[:valid_edges]))) + self.assertFalse(bool(np.any(graph.edge_mask[valid_edges:]))) + destination = graph.edge_index[1, :valid_edges] + self.assertTrue(bool(np.all(destination[:-1] <= destination[1:]))) + def test_oracle_set_equality_dense_periodic(self) -> None: """Periodic PBC: builder with pair_excl==(0,0) == builder() + apply.""" pe = self._pair_excl([(0, 0)]) diff --git a/source/tests/pt/model/test_descriptor_dpa1_triton.py b/source/tests/pt/model/test_descriptor_dpa1_triton.py new file mode 100644 index 0000000000..391c4921e9 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_dpa1_triton.py @@ -0,0 +1,593 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for the opt-in Triton inference kernel of the DPA1 descriptor +(``se_atten`` with ``attn_layer == 0``, in either ``strip`` or ``concat`` +tebd-input mode), enabled via the ``DP_TRITON_INFER`` level (see +:func:`deepmd.kernels.utils.triton_infer_level`). + +Three properties are covered: + +1. Numerical correctness of the fused environment convolution ``se_conv`` + (forward and the inference force backward) against the eager reference, + for both gate modes (``gated`` 1 = strip type-pair gate, 0 = concat no gate), + across every residual structure (doubling / identity / none), both inlined + activations (``tanh`` / ``silu``), both power-of-two and non-power-of-two + channel widths (the latter exercising the masked-padding path), and with or + without a per-layer timestep. +2. ``make_fx`` composability: the operator and its backward each trace as a + single opaque node, mirroring the ``pt_expt`` inference path that lowers a + ``make_fx`` graph through ``torch.compile`` / export. +3. Descriptor-level parity: routing ``DescrptDPA1`` through the fused kernel + reproduces the dense reference descriptor and its coordinate gradient for + both tebd-input modes across the doubling, identity, ``resnet_dt`` and + non-power-of-two embedding shapes, and the routed forward remains + ``torch.jit.script``-able (the LAMMPS freeze path prunes the operator and + keeps the reference). +""" + +import os +import unittest + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.kernels.triton.dpa1.activation import ( + TRITON_AVAILABLE, +) +from deepmd.kernels.triton.dpa1.edge_conv import ( + _edge_conv_reference, + _edge_conv_reference_backward, + edge_conv, +) +from deepmd.kernels.triton.dpa1.gemm_fp16x3 import ( + embed_gemm_fp16x3, +) +from deepmd.kernels.triton.dpa1.se_conv import ( + _se_conv_reference, + se_conv, +) +from deepmd.kernels.triton.dpa1.tile_configs import ( + DEFAULT_CONFIG, + resolve_conv_config, +) +from deepmd.pt.model.descriptor import ( + DescrptDPA1, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + +_CUDA = torch.cuda.is_available() +_GPU = unittest.skipUnless( + _CUDA and TRITON_AVAILABLE, + "CUDA and Triton are required for the DPA1 fused kernel", +) + + +def _rand_conv_inputs(nfnl, nnei, ng, resnet_mult, has_idt, ntype_pair, device, seed=0): + gen = torch.Generator(device=device).manual_seed(seed) + h1_dim = ng // resnet_mult if resnet_mult > 0 else ng + z2 = torch.randn(nfnl, nnei, ng, device=device, generator=gen) + h1 = torch.randn(nfnl, nnei, h1_dim, device=device, generator=gen) + idt = ( + 0.1 + torch.rand(ng, device=device, generator=gen) + if has_idt + else torch.ones(ng, device=device) + ) + tt = torch.randn(ntype_pair, ng, device=device, generator=gen) * 0.3 + idx = torch.randint(0, ntype_pair, (nfnl * nnei,), device=device, generator=gen) + sw = torch.rand(nfnl, nnei, device=device, generator=gen) + rr = torch.randn(nfnl, nnei, 4, device=device, generator=gen) + return z2, h1, idt, tt, idx, sw, rr + + +class TestSeConvConfig(unittest.TestCase): + """Launch-configuration resolution (CPU-safe, no kernel launch).""" + + def test_level1_returns_default(self) -> None: + self.assertEqual(resolve_conv_config(128, 64, level=1), DEFAULT_CONFIG) + + def test_level0_returns_default(self) -> None: + self.assertEqual(resolve_conv_config(128, 64, level=0), DEFAULT_CONFIG) + + def test_level2_unknown_shape_falls_back(self) -> None: + # An unswept channel width can only fall back to the universal default. + self.assertEqual(resolve_conv_config(777, 333, level=2), DEFAULT_CONFIG) + + +@_GPU +class TestSeConvOp(unittest.TestCase): + """Forward / backward correctness and make_fx composability of ``se_conv``. + + Every residual structure (``resnet_mult`` 2/1/0) is checked both with and + without a per-layer timestep. + """ + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + # ``ng = 128`` exercises the power-of-two fast path (no padding, and the + # doubling fold via a reshape); ``ng = 100`` exercises the masked-padding + # path (padded to 128) and the direct doubling fold at the true ``H1``. + # ``act`` selects the inlined activation: 0 = tanh, 1 = silu. + self.cases = [ + (ng, mult, has_idt, act, gated) + for ng in (128, 100) + for mult in (2, 1, 0) + for has_idt in (False, True) + for act in (0, 1) + for gated in (1, 0) + ] + + def test_forward_matches_reference(self) -> None: + for ng, mult, has_idt, act, gated in self.cases: + with self.subTest(ng=ng, mult=mult, has_idt=has_idt, act=act, gated=gated): + z2, h1, idt, tt, idx, sw, rr = _rand_conv_inputs( + 512, 47, ng, mult, has_idt, 169, self.device + ) + ref = _se_conv_reference(z2, h1, idt, tt, idx, sw, rr, mult, act, gated) + got = se_conv(z2, h1, idt, tt, idx, sw, rr, mult, act, gated) + rel = (got - ref).abs().max() / ref.abs().max() + self.assertLess(rel.item(), 1e-5) + + def test_backward_matches_reference(self) -> None: + for ng, mult, has_idt, act, gated in self.cases: + with self.subTest(ng=ng, mult=mult, has_idt=has_idt, act=act, gated=gated): + z2, h1, idt, tt, idx, sw, rr = _rand_conv_inputs( + 512, 47, ng, mult, has_idt, 169, self.device + ) + gout = torch.randn_like( + _se_conv_reference(z2, h1, idt, tt, idx, sw, rr, mult, act, gated) + ) + ref_in = [ + t.detach().clone().requires_grad_(True) for t in (z2, h1, sw, rr) + ] + _se_conv_reference( + ref_in[0], + ref_in[1], + idt, + tt, + idx, + ref_in[2], + ref_in[3], + mult, + act, + gated, + ).backward(gout) + got_in = [ + t.detach().clone().requires_grad_(True) for t in (z2, h1, sw, rr) + ] + se_conv( + got_in[0], + got_in[1], + idt, + tt, + idx, + got_in[2], + got_in[3], + mult, + act, + gated, + ).backward(gout) + for name, a, b in zip( + ("z2", "h1", "sw", "rr"), + (t.grad for t in ref_in), + (t.grad for t in got_in), + strict=True, + ): + # h1 carries no gradient without a residual last layer; + # sw carries none in concat (no gate). Both are None on the + # reference side, and the kernel returns an exact zero. + if (name == "h1" and mult == 0) or (name == "sw" and not gated): + self.assertLess(b.abs().max().item(), 1e-12) + continue + rel = (a - b).abs().max() / a.abs().max().clamp_min(1e-20) + self.assertLess(rel.item(), 1e-5, msg=f"grad {name}") + + def test_make_fx_force_trace(self) -> None: + z2, h1, idt, tt, idx, sw, rr = _rand_conv_inputs( + 512, 47, 128, 2, False, 169, self.device + ) + gout = torch.randn_like( + _se_conv_reference(z2, h1, idt, tt, idx, sw, rr, 2, 0, 1) + ) + + def force_fn(z2, h1, idt, tt, idx, sw, rr): + z2r = z2.detach().requires_grad_(True) + out = se_conv(z2r, h1, idt, tt, idx, sw, rr, 2, 0, 1) + (grad_z2,) = torch.autograd.grad(out, z2r, gout) + return out, grad_z2 + + traced = make_fx(force_fn)(z2, h1, idt, tt, idx, sw, rr) + n_fwd = sum(1 for n in traced.graph.nodes if "se_conv.default" in str(n.target)) + n_bwd = sum(1 for n in traced.graph.nodes if "se_conv_bwd" in str(n.target)) + self.assertEqual(n_fwd, 1) + self.assertEqual(n_bwd, 1) + out_ref, _ = force_fn(z2, h1, idt, tt, idx, sw, rr) + out_traced, _ = traced(z2, h1, idt, tt, idx, sw, rr) + torch.testing.assert_close(out_traced, out_ref) + + +def _rand_edge_inputs( + n_edge, n_node, ng, resnet_mult, has_idt, ntype_pair, device, seed=0 +): + gen = torch.Generator(device=device).manual_seed(seed) + h1_dim = ng // resnet_mult if resnet_mult > 0 else ng + z2 = torch.randn(n_edge, ng, device=device, generator=gen) + h1 = torch.randn(n_edge, h1_dim, device=device, generator=gen) + idt = ( + 0.1 + torch.rand(ng, device=device, generator=gen) + if has_idt + else torch.ones(ng, device=device) + ) + tt = torch.randn(ntype_pair, ng, device=device, generator=gen) * 0.3 + idx = torch.randint(0, ntype_pair, (n_edge,), device=device, generator=gen) + sw = torch.rand(n_edge, device=device, generator=gen) + rr = torch.randn(n_edge, 4, device=device, generator=gen) + dst = torch.randint(0, n_node, (n_edge,), device=device, generator=gen) + edge_mask = (torch.rand(n_edge, device=device, generator=gen) > 0.1).to( + torch.float32 + ) + return z2, h1, idt, tt, idx, sw, rr, dst, edge_mask + + +@_GPU +class TestEdgeConvOp(unittest.TestCase): + """Forward / backward correctness and make_fx composability of the graph + ``edge_conv`` (node-parallel CSR segment reduction), for both gate modes + (``gated`` 1 = strip type-pair gate, 0 = concat no gate). + """ + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + self.n_edge, self.n_node = 4096, 256 + # ng = 128 exercises the power-of-two fast path; ng = 100 the masked + # padding path; act 0 = tanh, 1 = silu; gated 1 = strip, 0 = concat. + self.cases = [ + (ng, mult, has_idt, act, gated) + for ng in (128, 100) + for mult in (2, 1, 0) + for has_idt in (False, True) + for act in (0, 1) + for gated in (1, 0) + ] + + def test_forward_matches_reference(self) -> None: + for ng, mult, has_idt, act, gated in self.cases: + with self.subTest(ng=ng, mult=mult, has_idt=has_idt, act=act, gated=gated): + z2, h1, idt, tt, idx, sw, rr, dst, em = _rand_edge_inputs( + self.n_edge, self.n_node, ng, mult, has_idt, 169, self.device + ) + args = ( + z2, + h1, + idt, + tt, + idx, + sw, + rr, + dst, + em, + self.n_node, + mult, + act, + gated, + ) + ref = _edge_conv_reference(*args) + got = edge_conv(*args) + rel = (got - ref).abs().max() / ref.abs().max() + self.assertLess(rel.item(), 1e-5) + + def test_backward_matches_reference(self) -> None: + for ng, mult, has_idt, act, gated in self.cases: + with self.subTest(ng=ng, mult=mult, has_idt=has_idt, act=act, gated=gated): + z2, h1, idt, tt, idx, sw, rr, dst, em = _rand_edge_inputs( + self.n_edge, self.n_node, ng, mult, has_idt, 169, self.device + ) + tail = (dst, em, self.n_node, mult, act, gated) + gout = torch.randn_like( + _edge_conv_reference(z2, h1, idt, tt, idx, sw, rr, *tail) + ) + ref_in = [ + t.detach().clone().requires_grad_(True) for t in (z2, h1, sw, rr) + ] + _edge_conv_reference( + ref_in[0], ref_in[1], idt, tt, idx, ref_in[2], ref_in[3], *tail + ).backward(gout) + got_in = [ + t.detach().clone().requires_grad_(True) for t in (z2, h1, sw, rr) + ] + edge_conv( + got_in[0], got_in[1], idt, tt, idx, got_in[2], got_in[3], *tail + ).backward(gout) + for name, a, b in zip( + ("z2", "h1", "sw", "rr"), + (t.grad for t in ref_in), + (t.grad for t in got_in), + strict=True, + ): + # h1 carries no gradient without a residual last layer; sw + # carries none in concat (no gate). Both are None on the + # reference side; the kernel returns an exact zero. + if (name == "h1" and mult == 0) or (name == "sw" and not gated): + self.assertLess(b.abs().max().item(), 1e-12) + continue + rel = (a - b).abs().max() / a.abs().max().clamp_min(1e-20) + self.assertLess(rel.item(), 1e-5, msg=f"grad {name}") + + def test_make_fx_force_trace(self) -> None: + # The strip gate (gated == 1) exercises the type-pair table gather inside + # the opaque operator; it must still trace as one fwd + one bwd node. + z2, h1, idt, tt, idx, sw, rr, dst, em = _rand_edge_inputs( + 2048, 128, 128, 2, False, 169, self.device + ) + + def force_fn(rr): + rrr = rr.detach().requires_grad_(True) + out = edge_conv(z2, h1, idt, tt, idx, sw, rrr, dst, em, 128, 2, 0, 1) + (grad_rr,) = torch.autograd.grad(out.sum(), rrr) + return out, grad_rr + + traced = make_fx(force_fn)(rr) + targets = [str(n.target) for n in traced.graph.nodes] + self.assertEqual(sum("edge_conv.default" in t for t in targets), 1) + self.assertEqual(sum("edge_conv_bwd" in t for t in targets), 1) + + +class TestEdgeConvReferenceBackward(unittest.TestCase): + """Closed-form CPU backward vs autograd through the eager reference. + + ``_edge_conv_reference_backward`` is the operator's non-CUDA gradient (used + when the pt_expt graph is traced or run off the GPU); it must be a closed + form so the fallback composes under ``make_fx``. This validates it against + ``torch.autograd`` through the eager forward across every residual structure + and both gate modes. CPU-only, so it runs without a GPU. + """ + + def _check(self, ng: int, mult: int, gated: int, act: int) -> None: + dev = torch.device("cpu") + n_node = 64 + z2, h1, idt, tt, idx, sw, rr, dst, em = _rand_edge_inputs( + 512, n_node, ng, mult, True, 37, dev + ) + tail = (dst, em, n_node, mult, act, gated) + gout = torch.randn_like( + _edge_conv_reference(z2, h1, idt, tt, idx, sw, rr, *tail) + ) + leaves = [t.detach().clone().requires_grad_(True) for t in (z2, h1, sw, rr)] + _edge_conv_reference( + leaves[0], leaves[1], idt, tt, idx, leaves[2], leaves[3], *tail + ).backward(gout) + closed = dict( + zip( + ("z2", "h1", "sw", "rr"), + _edge_conv_reference_backward( + gout, z2, h1, idt, tt, idx, sw, rr, *tail + ), + strict=True, + ) + ) + for name, leaf in zip(("z2", "h1", "sw", "rr"), leaves, strict=True): + c = closed[name] + if (name == "h1" and mult == 0) or (name == "sw" and not gated): + self.assertLess(c.abs().max().item(), 1e-12) + continue + rel = (leaf.grad - c).abs().max() / leaf.grad.abs().max().clamp_min(1e-20) + self.assertLess(rel.item(), 1e-6, msg=f"grad {name}") + + def test_closed_form_matches_autograd(self) -> None: + for mult in (2, 1, 0): + for gated in (1, 0): + for act in (0, 1): + with self.subTest(mult=mult, gated=gated, act=act): + self._check(128, mult, gated, act) + + +@_GPU +class TestEmbedGemmFp16x3(unittest.TestCase): + """fp16x3 embedding GEMM: fp32-reference accuracy (forward + input gradient) + and ``make_fx`` composability. Only the large-M embedding shapes are the + intended target; the operator stays numerically correct everywhere. + """ + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + + def test_matches_fp32_reference(self) -> None: + for m, k, n in [ + (65536, 64, 128), + (65536, 128, 64), + (65536, 32, 64), + (4096, 256, 256), + ]: + with self.subTest(m=m, k=k, n=n): + gen = torch.Generator(device=self.device).manual_seed(0) + a = ( + torch.randn(m, k, device=self.device, generator=gen) * 0.3 + ).requires_grad_(True) + b = torch.randn(k, n, device=self.device, generator=gen) * 0.3 + ref = a.detach().double() @ b.double() + c = embed_gemm_fp16x3(a, b) + fwd = (c.double() - ref).abs().max() / ref.abs().max() + self.assertLess(fwd.item(), 1e-5) + go = torch.randn_like(c) + (ga,) = torch.autograd.grad(c, a, go) + ga_ref = go.double() @ b.double().t() + bwd = (ga.double() - ga_ref).abs().max() / ga_ref.abs().max() + self.assertLess(bwd.item(), 1e-5, msg="grad_a") + + def test_make_fx_force_trace(self) -> None: + a = torch.randn(4096, 64, device=self.device) + b = torch.randn(64, 128, device=self.device) + go = torch.randn(4096, 128, device=self.device) + + def force_fn(a): + ar = a.detach().requires_grad_(True) + c = embed_gemm_fp16x3(ar, b) + (grad_a,) = torch.autograd.grad(c.sum(), ar, go.new_ones(())) + return c, grad_a + + traced = make_fx(force_fn)(a) + targets = [str(n.target) for n in traced.graph.nodes] + # forward + backward each trace as an opaque fp16x3 GEMM node. + self.assertGreaterEqual(sum("embed_gemm_fp16x3" in t for t in targets), 2) + + +def _build_dpa1( + device, + neuron, + resnet_dt=False, + activation_function="tanh", + tebd_input_mode="strip", +): + des = DescrptDPA1( + rcut=6.0, + rcut_smth=2.0, + sel=20, + ntypes=2, + neuron=neuron, + axis_neuron=4, + tebd_dim=8, + tebd_input_mode=tebd_input_mode, + attn=8, + attn_layer=0, + resnet_dt=resnet_dt, + activation_function=activation_function, + precision="float32", + seed=1, + ).to(device) + des.eval() + return des + + +@_GPU +class TestSeAttenRouting(unittest.TestCase): + """Descriptor-level parity and TorchScript safety of the fused route.""" + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + gen = torch.Generator(device=self.device).manual_seed(3) + n = 24 + self.coord = (torch.rand(1, n, 3, generator=gen, device=self.device) * 8.0).to( + env.GLOBAL_PT_FLOAT_PRECISION + ) + self.atype = torch.randint(0, 2, (1, n), generator=gen, device=self.device) + self.box = ( + (torch.eye(3, device=self.device) * 8.0) + .reshape(1, 9) + .to(env.GLOBAL_PT_FLOAT_PRECISION) + ) + + def _run(self, des, coord): + ec, ea, mapping, nlist = extend_input_and_build_neighbor_list( + coord, + self.atype, + des.get_rcut(), + des.get_sel(), + mixed_types=des.mixed_types(), + box=self.box, + ) + return des(ec, ea, nlist, mapping=mapping) + + def _assert_parity(self, des) -> None: + self.assertTrue(des.se_atten.se_conv_eligible) + block = des.se_atten + block.use_triton_infer = False + c0 = self.coord.detach().clone().requires_grad_(True) + d0 = self._run(des, c0)[0] + (g0,) = torch.autograd.grad(d0.sum(), c0) + block.use_triton_infer = True + c1 = self.coord.detach().clone().requires_grad_(True) + d1 = self._run(des, c1)[0] + (g1,) = torch.autograd.grad(d1.sum(), c1) + torch.testing.assert_close(d1, d0, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(g1, g0, atol=1e-5, rtol=1e-5) + + def test_parity_doubling(self) -> None: + self._assert_parity(_build_dpa1(self.device, [8, 16, 32])) + + def test_parity_identity(self) -> None: + self._assert_parity(_build_dpa1(self.device, [8, 16, 16])) + + def test_parity_resnet_dt(self) -> None: + self._assert_parity(_build_dpa1(self.device, [8, 16, 32], resnet_dt=True)) + + def test_parity_nonpow2_doubling(self) -> None: + # Non-power-of-two widths (ng = 50, H1 = 25) route through the kernel's + # masked-padding path and the direct (non-reshape) doubling fold. + self._assert_parity(_build_dpa1(self.device, [12, 25, 50])) + + def test_parity_nonpow2_identity(self) -> None: + self._assert_parity(_build_dpa1(self.device, [12, 25, 25])) + + def test_parity_silu(self) -> None: + self._assert_parity( + _build_dpa1(self.device, [8, 16, 32], activation_function="silu") + ) + + def test_parity_concat_doubling(self) -> None: + # Concat mode: the type feature enters the embedding input (no gate); + # ``se_conv`` runs with ``gated == 0``. + self._assert_parity( + _build_dpa1(self.device, [8, 16, 32], tebd_input_mode="concat") + ) + + def test_parity_concat_identity(self) -> None: + self._assert_parity( + _build_dpa1(self.device, [8, 16, 16], tebd_input_mode="concat") + ) + + def test_parity_concat_nonpow2(self) -> None: + self._assert_parity( + _build_dpa1(self.device, [12, 25, 50], tebd_input_mode="concat") + ) + + def test_parity_level3_fp16x3(self) -> None: + # DP_TRITON_INFER=3 folds the z2 embedding GEMM into fp16x3; the routed + # descriptor and its coordinate gradient must match the fp32 (level-2) + # output to ~fp32 rounding (~2^-22 step). + des = _build_dpa1(self.device, [8, 16, 32], tebd_input_mode="concat") + des.se_atten.use_triton_infer = True + + def run(level): + os.environ["DP_TRITON_INFER"] = level + c = self.coord.detach().clone().requires_grad_(True) + d = self._run(des, c)[0] + (g,) = torch.autograd.grad(d.sum(), c) + return d.detach(), g.detach() + + saved = os.environ.get("DP_TRITON_INFER") + try: + d2, g2 = run("2") + d3, g3 = run("3") + finally: + if saved is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved + torch.testing.assert_close(d3, d2, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(g3, g2, atol=1e-4, rtol=1e-4) + + def test_ineligible_unsupported_activation(self) -> None: + # The kernel inlines only ``tanh`` / ``silu``; any other last-layer + # activation must keep the dense reference path. + des = _build_dpa1(self.device, [8, 16, 32], activation_function="gelu") + self.assertFalse(des.se_atten.se_conv_eligible) + + def test_torchscript_prunes_fused_branch(self) -> None: + # ``torch.jit.script`` (the LAMMPS freeze path) must build even with the + # fused route enabled; the operator lives behind ``is_scripting``. + des = _build_dpa1(self.device, [8, 16, 32]) + des.se_atten.use_triton_infer = True + self.assertIsNotNone(torch.jit.script(des)) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_env_mat_triton.py b/source/tests/pt/model/test_env_mat_triton.py new file mode 100644 index 0000000000..ad4d5ed825 --- /dev/null +++ b/source/tests/pt/model/test_env_mat_triton.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for the fused Triton environment-matrix kernel. + +The kernel (:mod:`deepmd.kernels.triton.env_mat`) is a drop-in for the +descriptors' ``prod_env_mat`` front end under ``DP_TRITON_INFER >= 1`` on CUDA. +These tests check, against the eager reference path (level 0): + +* forward parity of ``(env_mat, diff, switch)`` in fp32 and fp64, for the full + and radial-only outputs and both smooth switches; +* force parity, i.e. the coordinate gradient produced by the registered + closed-form backward; +* ``NaN``-safety of the exponential switch backward in fp32 (the factored + ``-a e w`` overflows; the kernel uses the fused ``-a exp(xarg - e)`` form); +* composability under ``make_fx`` (the operator is captured as one opaque node). +""" + +import os +import unittest + +import torch + +from deepmd.kernels.triton.env_mat import ( + TRITON_AVAILABLE, + env_mat, +) +from deepmd.pt.model.descriptor.env_mat import ( + prod_env_mat, +) +from deepmd.pt.utils import ( + env, +) + +from ...seed import ( + GLOBAL_SEED, +) + +_SKIP = not (torch.cuda.is_available() and TRITON_AVAILABLE) + + +def _set_level(level: int) -> None: + os.environ["DP_TRITON_INFER"] = str(level) + + +def _make_system(dtype, seed, nf=2, nloc=48, nnei=64, nall=160, ntypes=3): + """Random extended system with a self-free, partially padded neighbor list.""" + g = torch.Generator(device=env.DEVICE).manual_seed(seed) + coord = torch.rand(nf, nall, 3, generator=g, device=env.DEVICE, dtype=dtype) * 9.0 + nlist = torch.randint(0, nall, (nf, nloc, nnei), generator=g, device=env.DEVICE) + center = torch.arange(nloc, device=env.DEVICE)[None, :, None] + drop = (torch.rand(nf, nloc, nnei, generator=g, device=env.DEVICE) < 0.2) | ( + nlist == center + ) + nlist = torch.where(drop, torch.full_like(nlist, -1), nlist) + atype = torch.randint(0, ntypes, (nf, nloc), generator=g, device=env.DEVICE) + return coord, nlist, atype, ntypes + + +@unittest.skipIf(_SKIP, "CUDA and Triton are required for the env_mat kernel") +class TestEnvMatTriton(unittest.TestCase): + def tearDown(self) -> None: + os.environ.pop("DP_TRITON_INFER", None) + + def _eager(self, coord, nlist, atype, mean, std, rcut, rsmth, radial, prot, ue): + _set_level(0) + c = coord.clone().reshape(coord.shape[0], -1).requires_grad_() + e, d, s = prod_env_mat( + c, + nlist, + atype, + mean, + std, + rcut, + rsmth, + radial_only=radial, + protection=prot, + use_exp_switch=ue, + ) + return c, e, d, s + + def _triton(self, coord, nlist, atype, mean, std, rcut, rsmth, radial, prot, ue): + _set_level(1) + c = coord.clone().requires_grad_() + e, d, s = env_mat( + c, + nlist, + atype, + mean, + std, + rcut, + rsmth, + radial_only=radial, + protection=prot, + use_exp_switch=ue, + ) + return c, e, d, s + + def test_forward_and_force_parity(self) -> None: + rcut, rsmth = 6.0, 2.0 # representable scalars -> fp64 stays exact + for dtype in (torch.float32, torch.float64): + for radial in (False, True): + for ue in (False, True): + with self.subTest(dtype=dtype, radial=radial, use_exp=ue): + ch = 1 if radial else 4 + coord, nlist, atype, nt = _make_system(dtype, GLOBAL_SEED) + nnei = nlist.shape[2] + gg = torch.Generator(device=env.DEVICE).manual_seed(1) + mean = torch.randn( + nt, nnei, ch, generator=gg, device=env.DEVICE, dtype=dtype + ) + std = 0.5 + torch.rand( + nt, nnei, ch, generator=gg, device=env.DEVICE, dtype=dtype + ) + ge = torch.randn( + *(coord.shape[0], nlist.shape[1], nnei, ch), + generator=gg, + device=env.DEVICE, + dtype=dtype, + ) + gd = torch.randn( + *(coord.shape[0], nlist.shape[1], nnei, 3), + generator=gg, + device=env.DEVICE, + dtype=dtype, + ) + gs = torch.randn( + *(coord.shape[0], nlist.shape[1], nnei, 1), + generator=gg, + device=env.DEVICE, + dtype=dtype, + ) + + c1, e1, d1, s1 = self._triton( + coord, nlist, atype, mean, std, rcut, rsmth, radial, 0.0, ue + ) + (g1,) = torch.autograd.grad( + (e1 * ge).sum() + (d1 * gd).sum() + (s1 * gs).sum(), c1 + ) + c0, e0, d0, s0 = self._eager( + coord, nlist, atype, mean, std, rcut, rsmth, radial, 0.0, ue + ) + (g0,) = torch.autograd.grad( + (e0 * ge).sum() + (d0 * gd).sum() + (s0 * gs).sum(), c0 + ) + + ftol = 1e-4 if dtype == torch.float32 else 1e-10 + gtol = 1e-3 if dtype == torch.float32 else 1e-8 + self.assertLess((e1 - e0).abs().max().item(), ftol) + self.assertLess((d1 - d0).abs().max().item(), ftol) + self.assertLess((s1 - s0).abs().max().item(), ftol) + gscale = g0.abs().max().item() + 1e-30 + self.assertLess( + (g1 - g0.reshape_as(g1)).abs().max().item() / gscale, gtol + ) + + def test_exp_switch_backward_is_finite_fp32(self) -> None: + # Extreme window (a = 20 / rcut_smth large): the factored switch + # derivative overflows in fp32; the kernel must stay finite. + _set_level(1) + coord, nlist, atype, nt = _make_system(torch.float32, GLOBAL_SEED) + nnei = nlist.shape[2] + mean = torch.zeros(nt, nnei, 4, device=env.DEVICE) + std = torch.ones(nt, nnei, 4, device=env.DEVICE) + c = coord.clone().requires_grad_() + e, d, s = env_mat(c, nlist, atype, mean, std, 6.0, 1.0, use_exp_switch=True) + (g,) = torch.autograd.grad(e.sum() + d.sum() + s.sum(), c) + self.assertTrue(torch.isfinite(g).all().item()) + + def test_make_fx_opaque_operator(self) -> None: + from torch.fx.experimental.proxy_tensor import ( + make_fx, + ) + + _set_level(1) + coord, nlist, atype, nt = _make_system(torch.float32, GLOBAL_SEED) + nnei = nlist.shape[2] + mean = torch.zeros(nt, nnei, 4, device=env.DEVICE) + std = torch.ones(nt, nnei, 4, device=env.DEVICE) + + def fn(c, nl, at, mn, sd): + e, d, s = env_mat(c, nl, at, mn, sd, 6.0, 2.0) + return e.sum() + d.sum() + s.sum() + + gm = make_fx(fn, tracing_mode="fake")(coord, nlist, atype, mean, std) + targets = [str(n.target) for n in gm.graph.nodes] + self.assertTrue(any("env_mat.default" in t for t in targets)) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py new file mode 100644 index 0000000000..97920c65d8 --- /dev/null +++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py @@ -0,0 +1,2016 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Graph-lower CUDA mega-kernel path of the pt_expt DPA1 model. + +At ``DP_CUDA_INFER >= 1`` the concat, attention-free graph lower routes +through three fused CUDA operator suites (see ``source/op/pt``): + +* ``deepmd::dpa1_graph_descriptor`` -- the descriptor mega kernels + (:mod:`deepmd.kernels.cuda.dpa1.graph_descriptor`); +* ``deepmd::graph_fitting`` -- the fused energy fitting network + (:mod:`deepmd.kernels.cuda.graph_fitting`); +* ``deepmd::edge_force_virial`` -- the fused force / virial assembly + (:mod:`deepmd.kernels.cuda.edge_force_virial`). + +Covered properties: + +1. Descriptor parity (concat): the fused route reproduces the dpmodel + reference ``call_graph`` (grrg, rot_mat and the ``edge_vec`` gradient) + across the supported widths, activations, one/two-side and ``resnet_dt``, + with an fp64 ``edge_vec`` leaf against an fp32 model (the graph ``.pt2`` + ABI). +2. Descriptor parity (strip): the dpmodel graph reference does not implement + strip, so values are checked against the dense reference ``call`` on the + same neighbor list and the ``edge_vec`` gradient against the operator's + CPU implementation (an independent autograd formulation). +3. Eligibility gating: unsupported configurations fall back, and the traced + graph contains the operator only at ``DP_CUDA_INFER >= 1``. +4. Fitting parity: values and the descriptor gradient of the fused fitting + match the dpmodel reference. +5. Scatter parity: force / atom-virial / per-frame virial of the fused + assembly match the array-API reference on a multi-frame graph, and the + CPU trace-time implementation matches the CUDA one. +""" + +import copy +import dataclasses +import os +import unittest + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + +_CUDA = torch.cuda.is_available() + + +def _cuda_ops_loaded() -> bool: + if not _CUDA: + return False + from deepmd.kernels.cuda.dpa1.graph_descriptor import ( + op_available, + ) + + return op_available() + + +_GPU = unittest.skipUnless( + _CUDA and _cuda_ops_loaded(), + "CUDA and the compiled deepmd op library are required", +) + + +class _CudaLevel: + """Context manager pinning ``DP_CUDA_INFER`` and restoring it on exit.""" + + def __init__(self, level: str) -> None: + self.level = level + self.saved: str | None = None + + def __enter__(self) -> None: + self.saved = os.environ.get("DP_CUDA_INFER") + os.environ["DP_CUDA_INFER"] = self.level + + def __exit__(self, *exc: object) -> None: + if self.saved is None: + os.environ.pop("DP_CUDA_INFER", None) + else: + os.environ["DP_CUDA_INFER"] = self.saved + + +def _build_dpa1_expt( + device, + neuron, + one_side=False, + act="tanh", + resnet_dt=False, + tebd_input_mode="concat", + smooth=True, + ntypes=2, + axis_neuron=4, + tebd_dim=8, +): + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + des = DescrptDPA1( + rcut=6.0, + rcut_smth=2.0, + sel=40, + ntypes=ntypes, + neuron=neuron, + axis_neuron=axis_neuron, + tebd_dim=tebd_dim, + tebd_input_mode=tebd_input_mode, + attn_layer=0, + type_one_side=one_side, + activation_function=act, + resnet_dt=resnet_dt, + smooth_type_embedding=smooth, + precision="float32", + seed=1, + ).to(device) + des.eval() + return des + + +@_GPU +class TestDpa1GraphCudaDescriptor(unittest.TestCase): + """Parity and routing of the fused descriptor mega kernels.""" + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + gen = torch.Generator(device=self.device).manual_seed(3) + n = 48 + self.coord = (torch.rand(1, n, 3, generator=gen, device=self.device) * 9.0).to( + torch.float64 + ) + self.atype = torch.randint(0, 2, (1, n), generator=gen, device=self.device) + self.box = ( + (torch.eye(3, device=self.device) * 9.0).reshape(1, 9).to(torch.float64) + ) + + def _graph(self, des): + graph, atype, _dense = self._graph_and_dense(des) + return graph, atype + + def _graph_and_dense(self, des): + from deepmd.dpmodel.utils.neighbor_graph import ( + from_dense_quartet, + ) + + ec, ea, mp, nl = extend_input_and_build_neighbor_list( + self.coord, + self.atype, + des.get_rcut(), + des.get_sel(), + mixed_types=des.mixed_types(), + box=self.box, + ) + graph = from_dense_quartet(ec, nl, mp, compact=True) + return graph, self.atype.reshape(-1).to(self.device), (ec, ea, nl) + + def _run(self, des, graph, atype, level): + """(grrg, rot_mat, d_edge_vec) through the public call_graph route.""" + with _CudaLevel(level): + ev = graph.edge_vec.detach().clone().requires_grad_(True) + g = dataclasses.replace(graph, edge_vec=ev) + tebd = des.type_embedding.call() + grrg, rot = des.call_graph(g, atype, type_embedding=tebd) + # A non-uniform cotangent exercises every output channel. + cot = torch.linspace(0.5, 1.5, grrg.numel(), device=grrg.device).reshape( + grrg.shape + ) + (gvec,) = torch.autograd.grad((grrg * cot).sum(), ev) + return grrg.detach(), rot.detach(), gvec.detach() + + def _assert_parity(self, des) -> None: + self.assertTrue(des._fused_eligible("cuda")) + graph, atype = self._graph(des) + d0, r0, g0 = self._run(des, graph, atype, "0") + d1, r1, g1 = self._run(des, graph, atype, "1") + # fp32 compute against the fp32 reference; the gradient returns in the + # fp64 leaf dtype on both routes. + self.assertEqual(g1.dtype, torch.float64) + torch.testing.assert_close(d1, d0, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(r1, r0, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(g1, g0.to(g1.dtype), atol=1e-5, rtol=1e-4) + + def test_parity_two_side_tanh(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [32, 64, 128])) + + def test_parity_one_side_silu_resnet_dt(self) -> None: + self._assert_parity( + _build_dpa1_expt( + self.device, + [32, 64, 128], + one_side=True, + act="silu", + resnet_dt=True, + ) + ) + + def test_parity_narrow_silu(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [16, 32, 64], act="silu")) + + def test_parity_narrow_one_side_resnet_dt(self) -> None: + self._assert_parity( + _build_dpa1_expt(self.device, [16, 32, 64], one_side=True, resnet_dt=True) + ) + + def test_parity_width8(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [8, 16, 32])) + + def test_parity_uniform_width(self) -> None: + # Identity residual shape on both upper layers. + self._assert_parity(_build_dpa1_expt(self.device, [64, 64, 64])) + + def test_parity_mixed_identity_doubling(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [16, 32, 32], act="silu")) + self._assert_parity(_build_dpa1_expt(self.device, [32, 32, 64])) + + def test_rotation_output_can_be_suppressed(self) -> None: + from deepmd.kernels.cuda.dpa1.graph_descriptor import ( + dpa1_graph_descriptor, + ) + + des = _build_dpa1_expt(self.device, [16, 32, 64]) + graph, atype = self._graph(des) + type_embedding = des.type_embedding.call() + descriptor, rotation = dpa1_graph_descriptor( + des, + graph, + atype, + type_embedding, + write_rotation=False, + ) + reference, _ = dpa1_graph_descriptor( + des, + graph, + atype, + type_embedding, + ) + self.assertEqual(rotation.shape, (0, des.se_atten.neuron[-1], 3)) + torch.testing.assert_close(descriptor, reference) + + def _assert_strip_parity(self, des) -> None: + """Strip mode: values vs the dense reference ``call`` on the same + neighbor list; the ``edge_vec`` gradient vs the operator's CPU + implementation (the dpmodel graph reference does not implement + strip). + """ + from deepmd.kernels.cuda.dpa1.graph_descriptor import ( + dpa1_graph_descriptor, + ) + + self.assertTrue(des._fused_eligible("cuda")) + graph, atype, (ec, ea, nl) = self._graph_and_dense(des) + nloc = self.atype.shape[1] + + # === Step 1. Values against the dense reference === + d_ref = des.call(ec, ea, nl)[0] # (1, nloc, out), fp64 output + tebd = des.type_embedding.call() + grrg, _rot = des._call_graph_cuda(graph, atype, tebd) + torch.testing.assert_close( + grrg[:nloc].to(torch.float64), d_ref[0], atol=1e-5, rtol=1e-5 + ) + + # === Step 2. edge_vec gradient against the CPU implementation === + def run(d, g, at, te): + ev = g.edge_vec.detach().clone().requires_grad_(True) + gg = dataclasses.replace(g, edge_vec=ev) + out, _ = dpa1_graph_descriptor(d, gg, at, te) + cot = torch.linspace(0.5, 1.5, out.numel(), device=out.device).reshape( + out.shape + ) + (gvec,) = torch.autograd.grad((out * cot).sum(), ev) + return gvec.detach() + + g_cuda = run(des, graph, atype, tebd) + des_cpu = copy.deepcopy(des).to("cpu") + graph_cpu = dataclasses.replace( + graph, + edge_vec=graph.edge_vec.cpu(), + edge_index=graph.edge_index.cpu(), + edge_mask=graph.edge_mask.cpu(), + destination_order=( + graph.destination_order.cpu() + if graph.destination_order is not None + else None + ), + destination_row_ptr=( + graph.destination_row_ptr.cpu() + if graph.destination_row_ptr is not None + else None + ), + source_order=( + graph.source_order.cpu() if graph.source_order is not None else None + ), + source_row_ptr=( + graph.source_row_ptr.cpu() if graph.source_row_ptr is not None else None + ), + n_node=graph.n_node.cpu(), + ) + g_cpu = run(des_cpu, graph_cpu, atype.cpu(), tebd.cpu()) + torch.testing.assert_close(g_cuda.cpu(), g_cpu, atol=1e-6, rtol=1e-4) + + def test_parity_strip_smooth_two_side(self) -> None: + self._assert_strip_parity( + _build_dpa1_expt(self.device, [32, 64, 128], tebd_input_mode="strip") + ) + + def test_parity_strip_smooth_one_side_silu(self) -> None: + self._assert_strip_parity( + _build_dpa1_expt( + self.device, + [16, 32, 64], + one_side=True, + act="silu", + tebd_input_mode="strip", + ) + ) + + def test_parity_strip_nosmooth(self) -> None: + self._assert_strip_parity( + _build_dpa1_expt( + self.device, [16, 32, 64], tebd_input_mode="strip", smooth=False + ) + ) + + def test_eligibility_gate(self) -> None: + # Non-power-of-two widths, shrinking layers and over-limit widths stay + # on the reference. Layers two and three support equal-or-doubling + # stacks; a first-layer residual remains on the reference path. + self.assertFalse( + _build_dpa1_expt(self.device, [24, 48, 96])._fused_eligible("cuda") + ) + self.assertFalse( + _build_dpa1_expt(self.device, [32, 64, 32])._fused_eligible("cuda") + ) + self.assertFalse( + _build_dpa1_expt(self.device, [64, 128, 128])._fused_eligible("cuda") + ) + self.assertTrue( + _build_dpa1_expt(self.device, [32, 64, 64])._fused_eligible("cuda") + ) + self.assertFalse( + _build_dpa1_expt( + self.device, + [8, 16, 32], + one_side=True, + tebd_dim=3, + )._fused_eligible("cuda") + ) + self.assertFalse( + _build_dpa1_expt(self.device, [32, 64, 64]) + .to(torch.float64) + ._fused_eligible("cuda") + ) + + def test_call_graph_routes_on_level(self) -> None: + des = _build_dpa1_expt(self.device, [32, 64, 128]) + tebd = des.type_embedding.call() + graph, atype = self._graph(des) + counts = {} + for level in ("0", "1"): + with _CudaLevel(level): + + def fn(edge_vec): + g = dataclasses.replace(graph, edge_vec=edge_vec) + return des.call_graph(g, atype, type_embedding=tebd) + + traced = make_fx(fn, tracing_mode="real")(graph.edge_vec) + counts[level] = sum( + "dpa1_graph_descriptor" in str(n.target) for n in traced.graph.nodes + ) + self.assertEqual(counts["0"], 0) + self.assertGreaterEqual(counts["1"], 1) + + +def _build_compressed_dpa1( + device, + neuron, + one_side=False, + act="tanh", + smooth=True, + min_nbor_dist=0.9, + ntypes=2, + axis_neuron=4, + tebd_dim=8, +): + """Strip DPA1 with the geometric embedding tabulated (``geo_compress``).""" + des = _build_dpa1_expt( + device, + neuron, + one_side=one_side, + act=act, + resnet_dt=False, + tebd_input_mode="strip", + smooth=smooth, + ntypes=ntypes, + axis_neuron=axis_neuron, + tebd_dim=tebd_dim, + ) + des.enable_compression(min_nbor_dist) + des.to(device) + des.eval() + return des + + +@_GPU +class TestDpa1GraphCudaCompress(unittest.TestCase): + """Parity and routing of the fused compressed (tabulated) descriptor kernel.""" + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + gen = torch.Generator(device=self.device).manual_seed(3) + n = 48 + self.coord = (torch.rand(1, n, 3, generator=gen, device=self.device) * 9.0).to( + torch.float64 + ) + self.atype = torch.randint(0, 2, (1, n), generator=gen, device=self.device) + self.box = ( + (torch.eye(3, device=self.device) * 9.0).reshape(1, 9).to(torch.float64) + ) + + def _graph_and_dense(self, des): + from deepmd.dpmodel.utils.neighbor_graph import ( + from_dense_quartet, + ) + + ec, ea, mp, nl = extend_input_and_build_neighbor_list( + self.coord, + self.atype, + des.get_rcut(), + des.get_sel(), + mixed_types=des.mixed_types(), + box=self.box, + ) + graph = from_dense_quartet(ec, nl, mp, compact=True, with_csr=True) + return graph, self.atype.reshape(-1).to(self.device), (ec, ea, nl) + + def _assert_parity(self, des) -> None: + """Values against the dense compressed reference ``call`` (which runs + ``tabulate_fusion_se_atten`` on the same table); the ``edge_vec`` + gradient against the operator's CPU implementation (an independent + autograd formulation through the same quintic table). + """ + from deepmd.kernels.cuda.dpa1.graph_compress import ( + dpa1_graph_compress, + ) + + self.assertTrue(des._fused_eligible("cuda")) + graph, atype, (ec, ea, nl) = self._graph_and_dense(des) + nloc = self.atype.shape[1] + + # === Step 1. Values against the dense tabulated reference === + d_ref = des.call(ec, ea, nl)[0] # (1, nloc, out), fp64 output + tebd = des.type_embedding.call() + with _CudaLevel("1"): + grrg, _rot = des._call_graph_cuda_compress(graph, atype, tebd) + torch.testing.assert_close( + grrg[:nloc].to(torch.float64), d_ref[0], atol=1e-5, rtol=1e-5 + ) + + # === Step 2. edge_vec gradient against the CPU implementation === + def run(d, g, at, te): + ev = g.edge_vec.detach().clone().requires_grad_(True) + gg = dataclasses.replace(g, edge_vec=ev) + out, _ = dpa1_graph_compress(d, gg, at, te) + cot = torch.linspace(0.5, 1.5, out.numel(), device=out.device).reshape( + out.shape + ) + (gvec,) = torch.autograd.grad((out * cot).sum(), ev) + return gvec.detach() + + g_cuda = run(des, graph, atype, tebd) + des_cpu = copy.deepcopy(des).to("cpu") + graph_cpu = dataclasses.replace( + graph, + edge_vec=graph.edge_vec.cpu(), + edge_index=graph.edge_index.cpu(), + edge_mask=graph.edge_mask.cpu(), + destination_order=graph.destination_order.cpu(), + destination_row_ptr=graph.destination_row_ptr.cpu(), + source_order=graph.source_order.cpu(), + source_row_ptr=graph.source_row_ptr.cpu(), + n_node=graph.n_node.cpu(), + ) + g_cpu = run(des_cpu, graph_cpu, atype.cpu(), tebd.cpu()) + torch.testing.assert_close(g_cuda.cpu(), g_cpu, atol=1e-6, rtol=1e-4) + + def test_parity_two_side_smooth(self) -> None: + # NG = 64: eight warps active in the moment backward. + self._assert_parity(_build_compressed_dpa1(self.device, [16, 32, 64])) + + def test_parity_wide_two_side(self) -> None: + # NG = 128: the moment backward covers the table in two channel blocks. + self._assert_parity(_build_compressed_dpa1(self.device, [32, 64, 128])) + + def test_parity_narrow_ng8(self) -> None: + # NG = 8: a single warp is active -- the shuffle fold must stay whole. + self._assert_parity(_build_compressed_dpa1(self.device, [16, 16, 8])) + + def test_parity_non_power_of_two_ng(self) -> None: + # NG = 100 (a non-power-of-two width): the kernel runs at the padded + # width 128 and the padding channels are sliced off. + self._assert_parity(_build_compressed_dpa1(self.device, [25, 50, 100])) + + def test_parity_one_side_silu(self) -> None: + self._assert_parity( + _build_compressed_dpa1(self.device, [16, 32, 64], one_side=True, act="silu") + ) + + def test_parity_nosmooth(self) -> None: + self._assert_parity( + _build_compressed_dpa1(self.device, [16, 32, 64], smooth=False) + ) + + def test_parity_axis16(self) -> None: + """Axis width 16 exercises the full symmetric Gram gradient.""" + self._assert_parity( + _build_compressed_dpa1(self.device, [16, 32, 64], axis_neuron=16) + ) + + def test_concat_type_embedding_wider_than_warp(self) -> None: + """The output tail covers every type-embedding channel.""" + self._assert_parity( + _build_compressed_dpa1(self.device, [16, 32, 64], tebd_dim=64) + ) + + def test_parity_four_types(self) -> None: + """Two-sided pair indexing covers a representative multi-element model.""" + self.atype = ( + torch.arange(self.atype.numel(), device=self.device).reshape( + self.atype.shape + ) + % 4 + ) + self._assert_parity(_build_compressed_dpa1(self.device, [16, 32, 64], ntypes=4)) + + def test_int32_edge_index_parity(self) -> None: + """Int32 edge addressing matches the default int64 graph ABI.""" + from deepmd.kernels.cuda.dpa1.graph_compress import ( + dpa1_graph_compress, + ) + + des = _build_compressed_dpa1(self.device, [16, 32, 64]) + graph64, atype, _ = self._graph_and_dense(des) + graph32 = dataclasses.replace( + graph64, + edge_index=graph64.edge_index.to(torch.int32), + destination_order=graph64.destination_order.to(torch.int32), + source_order=graph64.source_order.to(torch.int32), + ) + type_embedding = des.type_embedding.call() + + def run(graph) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) + graph = dataclasses.replace(graph, edge_vec=edge_vec) + descriptor, rotation = dpa1_graph_compress( + des, graph, atype, type_embedding + ) + cotangent = torch.linspace( + 0.5, 1.5, descriptor.numel(), device=descriptor.device + ).reshape(descriptor.shape) + (gradient,) = torch.autograd.grad((descriptor * cotangent).sum(), edge_vec) + return descriptor.detach(), rotation.detach(), gradient.detach() + + outputs64 = run(graph64) + outputs32 = run(graph32) + for output32, output64 in zip(outputs32, outputs64, strict=True): + torch.testing.assert_close(output32, output64, atol=1e-6, rtol=1e-6) + + def test_compact_canonical_descriptor_parity(self) -> None: + """Source-only topology matches the generic canonical operator.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + canonicalize_neighbor_graph, + ) + from deepmd.kernels.cuda.dpa1.canonical import ( + ensure_registered, + ) + from deepmd.kernels.cuda.dpa1.graph_compress import ( + dpa1_graph_compress, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + ) + + des = _build_compressed_dpa1(self.device, [16, 32, 64]) + graph, atype, _ = self._graph_and_dense(des) + graph = canonicalize_neighbor_graph( + dataclasses.replace(graph, n_local=graph.n_node), + atype.shape[0], + ) + compact = canonical_graph_from_neighbor_graph(graph) + type_embedding = des.type_embedding.call() + se = des.se_atten + inverse_stddev = torch.reciprocal(se.stddev[:, 0, :]).contiguous() + lower, upper, table_max, stride0, stride1 = ( + float(value) for value in des.compress_info[0].tolist()[:5] + ) + + generic_descriptor, _ = dpa1_graph_compress( + des, + graph, + atype, + type_embedding, + ) + ensure_registered() + descriptor, _rotation, moment = torch.ops.deepmd.dpa1_canonical_compress( + compact.edge_vec, + compact.source, + compact.destination_row_ptr, + atype, + type_embedding, + se.mean[:, 0, :].contiguous(), + inverse_stddev, + des.compress_data[0].contiguous(), + des.type_embd_data.contiguous(), + int(se.type_one_side), + int(des.concat_output_tebd), + 0, + int(se.smooth), + int(se.axis_neuron), + lower, + upper, + table_max, + stride0, + stride1, + float(se.rcut), + float(se.rcut_smth), + float(se.env_protection), + float(se.nnei), + ) + torch.testing.assert_close(descriptor, generic_descriptor) + + cotangent = torch.linspace( + 0.5, + 1.5, + descriptor.numel(), + dtype=descriptor.dtype, + device=descriptor.device, + ).reshape(descriptor.shape) + generic_edge_vec = graph.edge_vec.detach().requires_grad_(True) + generic_graph = dataclasses.replace(graph, edge_vec=generic_edge_vec) + generic_value, _ = dpa1_graph_compress( + des, + generic_graph, + atype, + type_embedding, + ) + (generic_gradient,) = torch.autograd.grad( + (generic_value * cotangent).sum(), + generic_edge_vec, + ) + compact_gradient = torch.ops.deepmd.dpa1_canonical_compress_backward( + cotangent, + None, + moment, + compact.edge_vec, + compact.source, + compact.destination_row_ptr, + atype, + se.mean[:, 0, :].contiguous(), + inverse_stddev, + des.compress_data[0].contiguous(), + des.type_embd_data.contiguous(), + int(se.type_one_side), + int(se.smooth), + int(se.axis_neuron), + lower, + upper, + table_max, + stride0, + stride1, + float(se.rcut), + float(se.rcut_smth), + float(se.env_protection), + float(se.nnei), + ) + physical_edge_count = int(compact.destination_row_ptr[-1].item()) + torch.testing.assert_close( + compact_gradient[:physical_edge_count], + generic_gradient[:physical_edge_count].to(compact_gradient.dtype), + atol=1e-6, + rtol=1e-5, + ) + torch.testing.assert_close( + compact_gradient[physical_edge_count:], + torch.zeros_like(compact_gradient[physical_edge_count:]), + ) + + def test_adaptive_resource_selection_large_graph(self) -> None: + """First-use tuning preserves the reference on a non-trivial graph.""" + generator = torch.Generator(device=self.device).manual_seed(37) + atom_count = 192 + box_length = 15.0 + self.coord = ( + torch.rand( + 1, + atom_count, + 3, + generator=generator, + device=self.device, + ) + * box_length + ).to(torch.float64) + self.atype = torch.randint( + 0, + 2, + (1, atom_count), + generator=generator, + device=self.device, + ) + self.box = ( + (torch.eye(3, device=self.device) * box_length) + .reshape(1, 9) + .to(torch.float64) + ) + self._assert_parity(_build_compressed_dpa1(self.device, [16, 32, 64])) + + def test_cached_topology_cutoff_mask_parity(self) -> None: + """Masked skin edges inside cached CSR rows contribute exactly zero.""" + des = _build_compressed_dpa1(self.device, [16, 32, 64]) + graph, atype, _ = self._graph_and_dense(des) + edge_mask = graph.edge_mask.clone() + self.assertTrue(bool(edge_mask[0])) + edge_mask[0] = False + graph = dataclasses.replace(graph, edge_mask=edge_mask) + type_embedding = des.type_embedding.call() + + def run(level: str) -> tuple[torch.Tensor, torch.Tensor]: + edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) + masked_graph = dataclasses.replace(graph, edge_vec=edge_vec) + with _CudaLevel(level): + descriptor, _ = des.call_graph( + masked_graph, atype, type_embedding=type_embedding + ) + cotangent = torch.linspace( + 0.5, 1.5, descriptor.numel(), device=descriptor.device + ).reshape(descriptor.shape) + (gradient,) = torch.autograd.grad((descriptor * cotangent).sum(), edge_vec) + return descriptor.detach(), gradient.detach() + + descriptor_ref, gradient_ref = run("0") + descriptor_fused, gradient_fused = run("1") + torch.testing.assert_close( + descriptor_fused, descriptor_ref, atol=1e-5, rtol=1e-5 + ) + torch.testing.assert_close(gradient_fused, gradient_ref, atol=1e-5, rtol=1e-4) + torch.testing.assert_close( + gradient_fused[0], torch.zeros_like(gradient_fused[0]) + ) + + def test_unsorted_edge_stream_uses_permutation_csr(self) -> None: + """The generic level-1 path preserves arbitrary edge payload order.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_edge_csr, + ) + + des = _build_compressed_dpa1(self.device, [16, 32, 64]) + graph, atype, _ = self._graph_and_dense(des) + generator = torch.Generator(device=self.device).manual_seed(19) + permutation = torch.randperm( + graph.edge_index.shape[1], generator=generator, device=self.device + ) + ( + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + graph.edge_index[:, permutation], + graph.edge_vec[permutation], + graph.edge_mask[permutation], + atype.shape[0], + ) + graph = dataclasses.replace( + graph, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=False, + ) + type_embedding = des.type_embedding.call() + + def run(level: str) -> tuple[torch.Tensor, torch.Tensor]: + current_edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) + current_graph = dataclasses.replace(graph, edge_vec=current_edge_vec) + with _CudaLevel(level): + descriptor, _ = des.call_graph( + current_graph, atype, type_embedding=type_embedding + ) + cotangent = torch.linspace( + 0.5, 1.5, descriptor.numel(), device=descriptor.device + ).reshape(descriptor.shape) + (gradient,) = torch.autograd.grad( + (descriptor * cotangent).sum(), current_edge_vec + ) + return descriptor.detach(), gradient.detach() + + descriptor_ref, gradient_ref = run("0") + descriptor_fused, gradient_fused = run("1") + torch.testing.assert_close( + descriptor_fused, descriptor_ref, atol=1e-5, rtol=1e-5 + ) + torch.testing.assert_close(gradient_fused, gradient_ref, atol=1e-5, rtol=1e-4) + + def test_level0_reference_parity(self) -> None: + """Level-0 graph path (original fused table op on the edge stream) + matches the dense compressed reference ``call``. + """ + des = _build_compressed_dpa1(self.device, [16, 32, 64]) + graph, atype, (ec, ea, nl) = self._graph_and_dense(des) + nloc = self.atype.shape[1] + d_ref = des.call(ec, ea, nl)[0] + tebd = des.type_embedding.call() + with _CudaLevel("0"): + grrg, _rot = des.call_graph(graph, atype, type_embedding=tebd) + torch.testing.assert_close( + grrg[:nloc].to(torch.float64), d_ref[0], atol=1e-5, rtol=1e-5 + ) + + def test_call_graph_routes_compress(self) -> None: + des = _build_compressed_dpa1(self.device, [16, 32, 64]) + tebd = des.type_embedding.call() + graph, atype, _ = self._graph_and_dense(des) + counts = {} + for level in ("0", "1"): + with _CudaLevel(level): + + def fn(edge_vec): + g = dataclasses.replace(graph, edge_vec=edge_vec) + return des.call_graph(g, atype, type_embedding=tebd) + + traced = make_fx(fn, tracing_mode="real")(graph.edge_vec) + counts[level] = sum( + "dpa1_graph_compress" in str(n.target) for n in traced.graph.nodes + ) + self.assertEqual(counts["0"], 0) + self.assertGreaterEqual(counts["1"], 1) + + def test_missing_csr_falls_back(self) -> None: + """A graph without optional CSR views remains a valid input.""" + des = _build_compressed_dpa1(self.device, [16, 32, 64]) + type_embedding = des.type_embedding.call() + graph, atype, _ = self._graph_and_dense(des) + graph = dataclasses.replace( + graph, + destination_order=None, + destination_row_ptr=None, + source_order=None, + source_row_ptr=None, + ) + with _CudaLevel("0"): + descriptor_ref, rotation_ref = des.call_graph( + graph, atype, type_embedding=type_embedding + ) + with _CudaLevel("1"): + descriptor, rotation = des.call_graph( + graph, atype, type_embedding=type_embedding + ) + torch.testing.assert_close(descriptor, descriptor_ref) + torch.testing.assert_close(rotation, rotation_ref) + + def test_empty_node_graph_has_zero_edge_gradient(self) -> None: + """An empty graph does not launch a zero-sized CUDA grid.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + + des = _build_compressed_dpa1(self.device, [16, 32, 64]) + edge_vec = torch.zeros( + 2, + 3, + dtype=torch.float64, + device=self.device, + requires_grad=True, + ) + order = torch.arange(2, dtype=torch.int64, device=self.device) + graph = NeighborGraph( + n_node=torch.zeros(1, dtype=torch.int64, device=self.device), + edge_index=torch.zeros(2, 2, dtype=torch.int64, device=self.device), + edge_vec=edge_vec, + edge_mask=torch.zeros(2, dtype=torch.bool, device=self.device), + destination_order=order, + destination_row_ptr=torch.zeros(1, dtype=torch.int64, device=self.device), + source_order=order, + source_row_ptr=torch.zeros(1, dtype=torch.int64, device=self.device), + destination_sorted=True, + ) + atype = torch.empty(0, dtype=torch.int64, device=self.device) + + with _CudaLevel("1"): + descriptor, rotation = des.call_graph( + graph, + atype, + type_embedding=des.type_embedding.call(), + ) + self.assertEqual(descriptor.shape, (0, des.get_dim_out())) + self.assertEqual(rotation.shape, (0, des.se_atten.neuron[-1], 3)) + (gradient,) = torch.autograd.grad(descriptor.sum(), edge_vec) + torch.testing.assert_close(gradient, torch.zeros_like(gradient)) + + def test_empty_node_cpu_operator_has_zero_edge_gradient(self) -> None: + """The registered CPU implementation preserves the empty-graph contract.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + + device = torch.device("cpu") + des = _build_compressed_dpa1(device, [16, 32, 64]) + edge_vec = torch.zeros(2, 3, dtype=torch.float64, requires_grad=True) + order = torch.arange(2, dtype=torch.int64) + graph = NeighborGraph( + n_node=torch.zeros(1, dtype=torch.int64), + edge_index=torch.zeros(2, 2, dtype=torch.int64), + edge_vec=edge_vec, + edge_mask=torch.zeros(2, dtype=torch.bool), + destination_order=order, + destination_row_ptr=torch.zeros(1, dtype=torch.int64), + source_order=order, + source_row_ptr=torch.zeros(1, dtype=torch.int64), + destination_sorted=True, + ) + atype = torch.empty(0, dtype=torch.int64) + + with _CudaLevel("1"): + descriptor, rotation = des.call_graph( + graph, + atype, + type_embedding=des.type_embedding.call(), + ) + self.assertEqual(descriptor.shape, (0, des.get_dim_out())) + self.assertEqual(rotation.shape, (0, des.se_atten.neuron[-1], 3)) + (gradient,) = torch.autograd.grad(descriptor.sum(), edge_vec) + torch.testing.assert_close(gradient, torch.zeros_like(gradient)) + + def test_torch_compile_first_backward(self) -> None: + """Inductor preserves the custom forward and its registered backward.""" + des = _build_compressed_dpa1(self.device, [16, 32, 64]) + type_embedding = des.type_embedding.call() + graph, atype, _ = self._graph_and_dense(des) + + def objective(edge_vec: torch.Tensor) -> torch.Tensor: + current_graph = dataclasses.replace(graph, edge_vec=edge_vec) + descriptor, _ = des.call_graph( + current_graph, atype, type_embedding=type_embedding + ) + return descriptor.square().sum() + + with _CudaLevel("1"): + traced = make_fx(objective, tracing_mode="real")(graph.edge_vec) + compiled = torch.compile(traced, fullgraph=True, dynamic=True) + + reference_input = graph.edge_vec.detach().clone().requires_grad_(True) + (reference_gradient,) = torch.autograd.grad( + traced(reference_input), reference_input + ) + compiled_input = graph.edge_vec.detach().clone().requires_grad_(True) + (compiled_gradient,) = torch.autograd.grad( + compiled(compiled_input), compiled_input + ) + torch.testing.assert_close( + compiled_gradient, reference_gradient, atol=1e-6, rtol=1e-5 + ) + + +@_GPU +class TestDpa1GraphCudaFitting(unittest.TestCase): + """Parity of the fused energy fitting network.""" + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + + def _build( + self, + resnet_dt, + act="tanh", + dim_descrpt=64, + neuron=None, + precision="float32", + ): + from deepmd.pt_expt.fitting.ener_fitting import ( + EnergyFittingNet, + ) + + if neuron is None: + neuron = [48, 48, 48] + fit = EnergyFittingNet( + ntypes=2, + dim_descrpt=dim_descrpt, + neuron=neuron, + resnet_dt=resnet_dt, + activation_function=act, + precision=precision, + mixed_types=True, + seed=2, + ).to(self.device) + fit.eval() + # A non-trivial per-type energy bias exercises the head epilogue. + fit.bias_atom_e = torch.tensor( + [[1.5], [-2.5]], dtype=torch.float64, device=self.device + ) + return fit + + def _assert_parity(self, fit) -> None: + from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + ) + + self.assertTrue(fitting_eligible(fit)) + gen = torch.Generator(device=self.device).manual_seed(5) + n = 37 + desc = torch.randn( + n, 64, generator=gen, device=self.device, dtype=torch.float32 + ) + atype = torch.randint(0, 2, (n,), generator=gen, device=self.device) + + def run(level): + with _CudaLevel(level): + x = desc.detach().clone().requires_grad_(True) + e = fit.call_graph(x, atype)["energy"] + (dx,) = torch.autograd.grad(e.sum(), x) + return e.detach(), dx.detach() + + e0, dx0 = run("0") + e1, dx1 = run("1") + self.assertEqual(e1.dtype, torch.float64) + torch.testing.assert_close(e1, e0, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(dx1, dx0, atol=1e-5, rtol=1e-5) + + 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_fparam_falls_back(self) -> None: + from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + ) + from deepmd.pt_expt.fitting.ener_fitting import ( + EnergyFittingNet, + ) + + fit = EnergyFittingNet( + ntypes=2, + dim_descrpt=64, + neuron=[48, 48], + numb_fparam=2, + precision="float32", + mixed_types=True, + seed=2, + ).to(self.device) + self.assertFalse(fitting_eligible(fit)) + + def test_width_doubling_residual_falls_back(self) -> None: + from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + ) + + fit = self._build( + resnet_dt=False, + dim_descrpt=32, + neuron=[64, 64], + ) + self.assertTrue(fit.nets[0].layers[0].resnet) + self.assertFalse(fitting_eligible(fit)) + + def test_float64_parameters_fall_back(self) -> None: + from deepmd.kernels.cuda.graph_fitting import ( + fitting_eligible, + ) + + self.assertFalse( + fitting_eligible(self._build(resnet_dt=False, precision="float64")) + ) + + +@_GPU +class TestDpa1GraphEnergyForce(unittest.TestCase): + """Parity of the fused end-to-end energy-force operator (DP_CUDA_INFER=2). + + The inference-only descriptor, fitting, descriptor-backward, and CSR scatter + sequence must reproduce the level-1 result with force and virial obtained + through ``autograd.grad``. + """ + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + gen = torch.Generator(device=self.device).manual_seed(3) + n = 48 + self.coord = (torch.rand(1, n, 3, generator=gen, device=self.device) * 9.0).to( + torch.float64 + ) + self.atype = torch.randint(0, 2, (1, n), generator=gen, device=self.device) + self.box = ( + (torch.eye(3, device=self.device) * 9.0).reshape(1, 9).to(torch.float64) + ) + + def _build_fitting(self, dim_descrpt): + from deepmd.pt_expt.fitting.ener_fitting import ( + EnergyFittingNet, + ) + + fit = EnergyFittingNet( + ntypes=2, + dim_descrpt=dim_descrpt, + neuron=[48, 48], + activation_function="silu", + precision="float32", + mixed_types=True, + seed=2, + ).to(self.device) + fit.eval() + fit.bias_atom_e = torch.tensor( + [[1.5], [-2.5]], dtype=torch.float64, device=self.device + ) + return fit + + def _graph(self, des): + from deepmd.dpmodel.utils.neighbor_graph import ( + from_dense_quartet, + ) + + ec, _ea, mp, nl = extend_input_and_build_neighbor_list( + self.coord, + self.atype, + des.get_rcut(), + des.get_sel(), + mixed_types=des.mixed_types(), + box=self.box, + ) + graph = from_dense_quartet(ec, nl, mp, compact=True, canonicalize=True) + return graph, self.atype.reshape(-1).to(self.device) + + def test_parity_vs_separate_ops(self) -> None: + from deepmd.kernels.cuda.dpa1.graph_energy_force import ( + dpa1_graph_energy_force, + ) + from deepmd.kernels.cuda.edge_force_virial import ( + edge_force_virial, + ) + + # A doubling stack exercises the retiled backward inside the fusion. + des = _build_dpa1_expt(self.device, [8, 16, 32], act="silu") + fit = self._build_fitting(des.get_dim_out()) + graph, atype = self._graph(des) + tebd = des.type_embedding.call() + n = atype.shape[0] + + with _CudaLevel("2"): + energy, atom_e, force, virial, atom_vir = dpa1_graph_energy_force( + des, + fit, + graph, + atype, + tebd, + torch.ones(n, dtype=torch.bool, device=self.device), + fit.bias_atom_e[:, 0].contiguous(), + node_capacity=n, + do_atomic_virial=True, + ) + + # Reference: the level-1 separate operators with the force from autograd. + with _CudaLevel("1"): + ev = graph.edge_vec.detach().clone().requires_grad_(True) + g2 = dataclasses.replace(graph, edge_vec=ev) + grrg, _rot = des.call_graph(g2, atype, tebd) + e_atom = fit.call_graph(grrg, atype)[fit.var_name] + (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( + g_e.to(force.dtype), + ev.detach().to(force.dtype), + graph.edge_index, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + graph.n_node, + n, + True, + ) + + self.assertEqual(energy.dtype, torch.float64) + torch.testing.assert_close( + atom_e, e_atom.to(torch.float64), atol=1e-6, rtol=1e-6 + ) + torch.testing.assert_close( + energy[0, 0], e_atom.sum().to(torch.float64), atol=1e-5, rtol=1e-6 + ) + torch.testing.assert_close(force, r_force, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(virial, r_virial, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(atom_vir, r_atom_vir, atol=1e-4, rtol=1e-4) + + def test_missing_csr_declines_energy_force_fusion(self) -> None: + """The caller can fall back when optional CSR views are absent.""" + for compressed in (False, True): + with self.subTest(compressed=compressed): + des = ( + _build_compressed_dpa1(self.device, [16, 32, 64]) + if compressed + else _build_dpa1_expt(self.device, [8, 16, 32], act="silu") + ) + fit = self._build_fitting(des.get_dim_out()) + graph, atype = self._graph(des) + graph = dataclasses.replace( + graph, + destination_order=None, + destination_row_ptr=None, + source_order=None, + source_row_ptr=None, + ) + result = des.fused_energy_force_graph( + fit, + graph, + atype, + torch.ones(atype.shape[0], dtype=torch.bool, device=self.device), + fit.bias_atom_e[:, 0].contiguous(), + do_atomic_virial=True, + ) + self.assertIsNone(result) + + def test_level2_reuses_virtual_and_pair_exclusion_masks(self) -> None: + from deepmd.pt_expt.model import ( + EnergyModel, + ) + + descriptor = _build_dpa1_expt(self.device, [8, 16, 32], act="silu") + fitting = self._build_fitting(descriptor.get_dim_out()) + model = EnergyModel( + descriptor, + fitting, + type_map=["A", "B"], + atom_exclude_types=[1], + pair_exclude_types=[(0, 1)], + ).to(self.device) + model.eval() + graph, atype = self._graph(descriptor) + atype = atype.clone() + atype[0] = -1 + args = ( + atype, + graph.n_node, + graph.n_node, + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + ) + + with _CudaLevel("1"): + reference = model.forward_common_lower_graph( + *args, + do_atomic_virial=True, + destination_sorted=graph.destination_sorted, + ) + with _CudaLevel("2"): + actual = model.forward_common_lower_graph( + *args, + do_atomic_virial=True, + destination_sorted=graph.destination_sorted, + ) + + self.assertEqual(set(actual), set(reference)) + for key in actual: + torch.testing.assert_close( + actual[key], + reference[key], + atol=1e-4, + rtol=1e-4, + ) + + def test_fused_energy_uses_owned_nodes_only(self) -> None: + from deepmd.kernels.cuda.edge_force_virial import ( + edge_force_virial, + ) + + des = _build_dpa1_expt(self.device, [8, 16, 32], act="silu") + fit = self._build_fitting(des.get_dim_out()) + graph, atype = self._graph(des) + n_node = atype.shape[0] + ownership = torch.arange(n_node, device=self.device) < n_node // 2 + output_bias = torch.tensor([2.0, -3.0], dtype=torch.float64, device=self.device) + + with _CudaLevel("2"): + fused = des.fused_energy_force_graph( + fit, + graph, + atype, + ownership, + fit.bias_atom_e[:, 0] + output_bias, + do_atomic_virial=True, + ) + assert fused is not None + energy, atom_energy, force, virial, atom_virial = fused + + with _CudaLevel("1"): + edge_vec = graph.edge_vec.detach().clone().requires_grad_(True) + current_graph = dataclasses.replace(graph, edge_vec=edge_vec) + descriptor, _ = des.call_graph( + current_graph, + atype, + type_embedding=des.type_embedding.call(), + ) + atom_energy_raw = fit.call_graph(descriptor, atype)[fit.var_name] + atom_energy_ref = (atom_energy_raw + output_bias[atype, None]) * ownership[ + :, None + ] + (edge_gradient,) = torch.autograd.grad(atom_energy_ref.sum(), edge_vec) + force_ref, atom_virial_ref, virial_ref = edge_force_virial( + edge_gradient.to(force.dtype), + edge_vec.detach().to(force.dtype), + graph.edge_index, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + graph.n_node, + n_node, + True, + ) + + torch.testing.assert_close(atom_energy, atom_energy_ref) + torch.testing.assert_close(energy[0], atom_energy_ref.sum(dim=0)) + torch.testing.assert_close(force, force_ref, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(virial, virial_ref, atol=1e-4, rtol=1e-4) + torch.testing.assert_close( + atom_virial, + atom_virial_ref, + atol=1e-4, + rtol=1e-4, + ) + + +@_GPU +class TestDpa1GraphCompressEnergyForce(unittest.TestCase): + """Parity of the fused compressed end-to-end energy-force operator. + + The tabulated descriptor, fitting, descriptor-backward, and CSR scatter + sequence must reproduce the level-1 result with force and virial obtained + through ``autograd.grad``. + """ + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + gen = torch.Generator(device=self.device).manual_seed(3) + n = 48 + self.coord = (torch.rand(1, n, 3, generator=gen, device=self.device) * 9.0).to( + torch.float64 + ) + self.atype = torch.randint(0, 2, (1, n), generator=gen, device=self.device) + self.box = ( + (torch.eye(3, device=self.device) * 9.0).reshape(1, 9).to(torch.float64) + ) + + def _build_fitting(self, dim_descrpt, ntypes=2): + from deepmd.pt_expt.fitting.ener_fitting import ( + EnergyFittingNet, + ) + + fit = EnergyFittingNet( + ntypes=ntypes, + dim_descrpt=dim_descrpt, + neuron=[64, 64, 64], + resnet_dt=True, + activation_function="silu", + precision="float32", + mixed_types=True, + seed=2, + ).to(self.device) + fit.eval() + fit.bias_atom_e = torch.linspace( + -2.5, 1.5, ntypes, dtype=torch.float64, device=self.device + ).reshape(ntypes, 1) + return fit + + def _graph(self, des): + from deepmd.dpmodel.utils.neighbor_graph import ( + from_dense_quartet, + ) + + ec, _ea, mp, nl = extend_input_and_build_neighbor_list( + self.coord, + self.atype, + des.get_rcut(), + des.get_sel(), + mixed_types=des.mixed_types(), + box=self.box, + ) + graph = from_dense_quartet(ec, nl, mp, compact=True, canonicalize=True) + return graph, self.atype.reshape(-1).to(self.device) + + def test_parity_vs_separate_ops(self) -> None: + from deepmd.kernels.cuda.dpa1.graph_compress import ( + dpa1_graph_compress, + dpa1_graph_compress_energy_force, + mega_eligible, + ) + from deepmd.kernels.cuda.edge_force_virial import ( + edge_force_virial, + ) + + self.atype = ( + torch.arange(self.atype.numel(), device=self.device).reshape( + self.atype.shape + ) + % 4 + ) + des = _build_compressed_dpa1( + self.device, + [16, 32, 64], + act="silu", + ntypes=4, + axis_neuron=16, + ) + self.assertTrue(mega_eligible(des)) + fit = self._build_fitting(des.get_dim_out(), ntypes=4) + graph, atype = self._graph(des) + edge_mask = graph.edge_mask.clone() + edge_mask[0] = False + graph = dataclasses.replace(graph, edge_mask=edge_mask) + tebd = des.type_embedding.call() + n = atype.shape[0] + + with _CudaLevel("2"): + energy, atom_e, force, virial, atom_vir = dpa1_graph_compress_energy_force( + des, + fit, + graph, + atype, + tebd, + torch.ones(n, dtype=torch.bool, device=self.device), + fit.bias_atom_e[:, 0].contiguous(), + node_capacity=n, + do_atomic_virial=True, + ) + + # Reference: the level-1 tabulated operator with the force from autograd. + with _CudaLevel("1"): + ev = graph.edge_vec.detach().clone().requires_grad_(True) + g2 = dataclasses.replace(graph, edge_vec=ev) + 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( + g_e.to(force.dtype), + ev.detach().to(force.dtype), + graph.edge_index, + graph.edge_mask, + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + graph.n_node, + n, + True, + ) + + self.assertEqual(energy.dtype, torch.float64) + torch.testing.assert_close( + atom_e, e_atom.to(torch.float64), atol=1e-6, rtol=1e-6 + ) + torch.testing.assert_close( + energy[0, 0], e_atom.sum().to(torch.float64), atol=1e-5, rtol=1e-6 + ) + torch.testing.assert_close(force, r_force, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(virial, r_virial, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(atom_vir, r_atom_vir, atol=1e-4, rtol=1e-4) + + def test_compact_canonical_model_trace(self) -> None: + """The eight-tensor deployment forward composes under symbolic make_fx.""" + from deepmd.pt_expt.model import ( + EnergyModel, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + ) + + descriptor = _build_compressed_dpa1( + self.device, + [16, 32, 64], + act="silu", + axis_neuron=16, + ) + fitting = self._build_fitting(descriptor.get_dim_out()) + model = EnergyModel( + descriptor, + fitting, + type_map=["A", "B"], + ).to(self.device) + model.eval() + graph, atype = self._graph(descriptor) + graph = dataclasses.replace(graph, n_local=graph.n_node) + compact = canonical_graph_from_neighbor_graph(graph) + inputs = ( + atype, + compact.n_node, + compact.n_local, + compact.source, + compact.edge_vec, + compact.destination_row_ptr, + compact.source_row_ptr, + compact.source_order, + ) + + reference = model.forward_lower_canonical_graph( + *inputs, + do_atomic_virial=True, + ) + traced = model.forward_lower_canonical_graph_exportable( + *inputs, + do_atomic_virial=True, + tracing_mode="real", + _allow_non_fake_inputs=True, + ) + actual = traced(*inputs) + self.assertEqual(set(actual), set(reference)) + for key in actual: + torch.testing.assert_close(actual[key], reference[key]) + + def test_compact_canonical_torch_export_contract(self) -> None: + """torch.export records the fixed eight-tensor deployment ABI.""" + from deepmd.pt_expt.model import ( + EnergyModel, + ) + from deepmd.pt_expt.utils.serialization import ( + _trace_and_export, + ) + + descriptor = _build_compressed_dpa1( + self.device, + [16, 32, 64], + act="silu", + axis_neuron=16, + ) + fitting = self._build_fitting(descriptor.get_dim_out()) + model = EnergyModel( + descriptor, + fitting, + type_map=["A", "B"], + ).to(self.device) + model.eval() + exported, metadata, _model_json, output_keys = _trace_and_export( + {"model": model.serialize()}, + do_atomic_virial=True, + lower_kind="dpa1_canonical", + ) + + self.assertEqual(metadata["lower_input_kind"], "dpa1_canonical") + self.assertEqual(metadata["graph_edge_dtype"], "float32") + self.assertNotIn("graph_index_dtype", metadata) + self.assertEqual( + output_keys, + ["atom_energy", "energy", "force", "virial", "mask", "atom_virial"], + ) + user_inputs = [ + spec + for spec in exported.graph_signature.input_specs + if spec.kind.name == "USER_INPUT" + ] + self.assertEqual(len(user_inputs), 8) + + def test_level2_permutation_csr_parity(self) -> None: + """Level 2 preserves force and virial for an arbitrary edge stream.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_edge_csr, + ) + from deepmd.kernels.cuda.dpa1.graph_compress import ( + dpa1_graph_compress_energy_force, + ) + + des = _build_compressed_dpa1(self.device, [16, 32, 64], act="silu") + fit = self._build_fitting(des.get_dim_out()) + canonical_graph, atype = self._graph(des) + edge_mask = canonical_graph.edge_mask.clone() + edge_mask[0] = False + canonical_graph = dataclasses.replace( + canonical_graph, + edge_mask=edge_mask, + ) + generator = torch.Generator(device=self.device).manual_seed(29) + permutation = torch.randperm( + canonical_graph.edge_index.shape[1], + generator=generator, + device=self.device, + ) + ( + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + canonical_graph.edge_index[:, permutation], + canonical_graph.edge_vec[permutation], + canonical_graph.edge_mask[permutation], + atype.shape[0], + ) + permutation_graph = dataclasses.replace( + canonical_graph, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=False, + ) + type_embedding = des.type_embedding.call() + + with _CudaLevel("2"): + canonical_output = dpa1_graph_compress_energy_force( + des, + fit, + canonical_graph, + atype, + type_embedding, + torch.ones(atype.shape[0], dtype=torch.bool, device=self.device), + fit.bias_atom_e[:, 0].contiguous(), + node_capacity=atype.shape[0], + do_atomic_virial=True, + ) + permutation_output = dpa1_graph_compress_energy_force( + des, + fit, + permutation_graph, + atype, + type_embedding, + torch.ones(atype.shape[0], dtype=torch.bool, device=self.device), + fit.bias_atom_e[:, 0].contiguous(), + node_capacity=atype.shape[0], + do_atomic_virial=True, + ) + for actual, expected in zip( + permutation_output, + canonical_output, + strict=True, + ): + torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4) + + +@_GPU +class TestEdgeForceVirialCuda(unittest.TestCase): + """Parity of the fused force / virial scatter (CUDA and CPU impls).""" + + def test_linear_csr_builder(self) -> None: + source = torch.tensor([2, 0, 1, 2, 0, 0, 0], dtype=torch.int64, device="cuda") + destination = torch.tensor( + [0, 0, 1, 2, 2, 0, 0], dtype=torch.int64, device="cuda" + ) + edge_index = torch.stack([source, destination]) + destination_order, destination_row_ptr, source_order, source_row_ptr = ( + torch.ops.deepmd.build_graph_csr(edge_index, 3, 5) + ) + + torch.testing.assert_close( + destination_order, torch.arange(7, dtype=torch.int64, device="cuda") + ) + torch.testing.assert_close( + destination_row_ptr, + torch.tensor([0, 2, 3, 5], dtype=torch.int64, device="cuda"), + ) + torch.testing.assert_close( + source_row_ptr, + torch.tensor([0, 2, 3, 5], dtype=torch.int64, device="cuda"), + ) + ordered_source = source[source_order[:5]] + torch.testing.assert_close( + ordered_source, + torch.tensor([0, 0, 1, 2, 2], dtype=torch.int64, device="cuda"), + ) + torch.testing.assert_close( + source_order[5:], + torch.tensor([5, 6], dtype=torch.int64, device="cuda"), + ) + + def _random_graph(self, device, n_frame=2, n_per_frame=9, n_edge=200): + gen = torch.Generator(device="cpu").manual_seed(7) + n_node = torch.full((n_frame,), n_per_frame, dtype=torch.int64) + total = int(n_node.sum()) + # Edges within each frame's node range, with a masked padding tail. + src, dst = [], [] + for f in range(n_frame): + lo = f * n_per_frame + src.append( + torch.randint(lo, lo + n_per_frame, (n_edge // n_frame,), generator=gen) + ) + dst.append( + torch.randint(lo, lo + n_per_frame, (n_edge // n_frame,), generator=gen) + ) + edge_index = torch.stack([torch.cat(src), torch.cat(dst)]) + E = edge_index.shape[1] + mask = torch.rand(E, generator=gen) > 0.15 + g_e = torch.randn(E, 3, generator=gen, dtype=torch.float64) + edge_vec = torch.randn(E, 3, generator=gen, dtype=torch.float64) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_edge_csr, + ) + + ( + edge_index, + payload, + _topology_mask, + dst_order, + dst_row_ptr, + src_order, + src_row_ptr, + ) = build_edge_csr( + edge_index, + torch.cat([g_e, edge_vec], dim=1), + torch.ones_like(mask), + total, + ) + g_e, edge_vec = payload[:, :3], payload[:, 3:] + move = lambda t: t.to(device) # noqa: E731 + return ( + 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), + total, + ) + + def _reference( + self, + g_e, + edge_vec, + edge_index, + mask, + dst_order, + dst_row_ptr, + src_order, + src_row_ptr, + n_node, + total, + ): + from deepmd.dpmodel.utils.neighbor_graph import ( + edge_force_virial, + ) + + return edge_force_virial( + g_e, edge_vec, edge_index, mask, n_node, node_capacity=total + ) + + def _fused( + self, + g_e, + edge_vec, + edge_index, + mask, + dst_order, + dst_row_ptr, + src_order, + src_row_ptr, + n_node, + total, + ): + from deepmd.kernels.cuda.edge_force_virial import ( + edge_force_virial, + ) + + return edge_force_virial( + g_e, + edge_vec, + edge_index, + mask, + dst_order, + dst_row_ptr, + src_order, + src_row_ptr, + n_node, + total, + True, + ) + + def _assert_device_parity(self, device) -> None: + args = self._random_graph(torch.device(device)) + f0, av0, v0 = self._reference(*args) + f1, av1, v1 = self._fused(*args) + torch.testing.assert_close(f1, f0, atol=1e-12, rtol=1e-12) + torch.testing.assert_close(av1, av0, atol=1e-12, rtol=1e-12) + torch.testing.assert_close(v1, v0, atol=1e-12, rtol=1e-12) + + def test_parity_cuda(self) -> None: + self._assert_device_parity("cuda") + + def test_compact_canonical_parity(self) -> None: + """Source-only force and virial match the generic dual-CSR operator.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + build_edge_csr, + ) + from deepmd.kernels.cuda.edge_force_virial import ( + canonical_edge_force_virial, + edge_force_virial, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + ) + + ( + g_e, + edge_vec, + edge_index, + _mask, + _dst_order, + _dst_row_ptr, + _src_row_ptr, + _src_order, + n_node, + total, + ) = self._random_graph(torch.device("cuda")) + ( + edge_index, + payload, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + edge_index, + torch.cat((g_e, edge_vec), dim=1), + torch.ones(edge_index.shape[1], dtype=torch.bool, device="cuda"), + total, + canonicalize=True, + ) + g_e = payload[:, :3].to(torch.float32) + edge_vec = payload[:, 3:].to(torch.float32) + graph = NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + n_local=n_node, + destination_order=destination_order, + destination_row_ptr=destination_row_ptr, + source_order=source_order, + source_row_ptr=source_row_ptr, + destination_sorted=True, + ) + compact = canonical_graph_from_neighbor_graph(graph) + + generic = edge_force_virial( + g_e, + edge_vec, + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + n_node, + total, + True, + ) + canonical = canonical_edge_force_virial( + g_e, + compact.edge_vec, + compact.destination_row_ptr, + compact.source_row_ptr, + compact.source_order, + compact.n_node, + 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_many_small_frames(self) -> None: + """Frame reduction is valid beyond the CUDA grid-y limit.""" + frame_count = 8192 + node = torch.arange(frame_count, device="cuda", dtype=torch.int64) + edge_index = torch.stack([node, node]) + edge_mask = torch.ones(frame_count, device="cuda", dtype=torch.bool) + generator = torch.Generator(device="cuda").manual_seed(23) + edge_gradient = torch.randn( + frame_count, + 3, + generator=generator, + device="cuda", + dtype=torch.float64, + ) + edge_vec = torch.randn( + frame_count, + 3, + generator=generator, + device="cuda", + dtype=torch.float64, + ) + order = node.clone() + row_ptr = torch.arange(frame_count + 1, device="cuda", dtype=torch.int64) + n_node_per_frame = torch.ones(frame_count, device="cuda", dtype=torch.int64) + + force, atom_virial, virial = self._fused( + edge_gradient, + edge_vec, + edge_index, + edge_mask, + order, + row_ptr, + order, + row_ptr, + n_node_per_frame, + frame_count, + ) + + expected_atom_virial = -torch.einsum("ei,ej->eij", edge_gradient, edge_vec) + torch.testing.assert_close(force, torch.zeros_like(force)) + torch.testing.assert_close(atom_virial, expected_atom_virial) + torch.testing.assert_close(virial, expected_atom_virial) + + def test_ragged_many_frame_parity(self) -> None: + """CSR scatter preserves force and virial across heterogeneous frames.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_edge_csr, + ) + + n_node = torch.tensor([1, 4, 2, 7] * 1024, dtype=torch.int64) + offsets = torch.cat( + [torch.zeros(1, dtype=torch.int64), torch.cumsum(n_node, dim=0)] + ) + source_parts = [] + destination_parts = [] + for frame in range(n_node.numel()): + nodes = torch.arange(int(offsets[frame]), int(offsets[frame + 1])) + source_parts.extend((nodes, torch.roll(nodes, shifts=1))) + destination_parts.extend((nodes, nodes)) + edge_index = torch.stack( + [torch.cat(source_parts), torch.cat(destination_parts)] + ) + generator = torch.Generator(device="cpu").manual_seed(31) + edge_mask = ( + torch.rand( + edge_index.shape[1], + generator=generator, + ) + > 0.2 + ) + edge_gradient = torch.randn( + edge_index.shape[1], + 3, + generator=generator, + dtype=torch.float64, + ) + edge_vec = torch.randn( + edge_index.shape[1], + 3, + generator=generator, + dtype=torch.float64, + ) + total = int(n_node.sum()) + ( + edge_index, + payload, + _topology_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + edge_index, + torch.cat([edge_gradient, edge_vec], dim=1), + torch.ones_like(edge_mask), + total, + ) + edge_gradient, edge_vec = payload[:, :3], payload[:, 3:] + args = ( + edge_gradient.cuda(), + edge_vec.cuda(), + edge_index.cuda(), + edge_mask.cuda(), + destination_order.cuda(), + destination_row_ptr.cuda(), + source_order.cuda(), + source_row_ptr.cuda(), + n_node.cuda(), + total, + ) + + reference = self._reference(*args) + fused = self._fused(*args) + for actual, expected in zip(fused, reference, strict=True): + torch.testing.assert_close(actual, expected, atol=1e-10, rtol=1e-10) + + def test_parity_cpu_impl(self) -> None: + # The CPU registration serves trace-time sample evaluation. + self._assert_device_parity("cpu") + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/descriptor/test_dpa1_triton.py b/source/tests/pt_expt/descriptor/test_dpa1_triton.py new file mode 100644 index 0000000000..f3dbdea931 --- /dev/null +++ b/source/tests/pt_expt/descriptor/test_dpa1_triton.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Graph-lower Triton path of the pt_expt DPA1 descriptor. + +The graph lower (``--lower-kind graph``) represents the neighbor list as a flat +edge stream. When ``DP_TRITON_INFER >= 1`` the attention-free block routes +``call_graph`` through the fused edge-parallel +:func:`~deepmd.kernels.triton.dpa1.edge_conv.edge_conv` operator, for both +tebd-input modes: ``concat`` (the type feature enters the embedding input, no +gate) and ``strip`` (the type feature factorizes into the type-pair gate +``gg = gg_s * (1 + tt[idx] * sw)``, fed by the per-edge switch from +``edge_env_mat(return_sw=True)``). + +Two properties are covered: + +1. Parity: the fused ``_call_graph_triton`` reproduces the dpmodel reference + ``call_graph`` (descriptor grrg, rotation matrix and the edge_vec force + gradient) across residual structures, one/two-side, activation and + non-power-of-two widths, in both tebd-input modes (strip with the smooth + switch on and off). +2. ``make_fx`` bake: ``edge_conv`` traces as a single opaque node (the pt_expt + graph-form ``.pt2`` export target), and ``call_graph`` routes to it only at + ``DP_TRITON_INFER >= 1``. +""" + +import dataclasses +import os +import unittest + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.kernels.triton.dpa1.activation import ( + TRITON_AVAILABLE, +) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + +_CUDA = torch.cuda.is_available() +_GPU = unittest.skipUnless( + _CUDA and TRITON_AVAILABLE, + "CUDA and Triton are required for the DPA1 graph edge kernel", +) + + +def _build_dpa1_expt( + device, + neuron, + one_side=False, + act="tanh", + precision="float64", + tebd_input_mode="concat", + smooth=True, +): + from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, + ) + + des = DescrptDPA1( + rcut=6.0, + rcut_smth=2.0, + sel=40, + ntypes=2, + neuron=neuron, + axis_neuron=4, + tebd_dim=8, + tebd_input_mode=tebd_input_mode, + attn_layer=0, + type_one_side=one_side, + activation_function=act, + precision=precision, + smooth_type_embedding=smooth, + seed=1, + ).to(device) + des.eval() + return des + + +@_GPU +class TestDpa1GraphRouting(unittest.TestCase): + """Graph-lower parity and make_fx bake of the fused ``edge_conv`` route.""" + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + gen = torch.Generator(device=self.device).manual_seed(3) + n = 24 + self.coord = (torch.rand(1, n, 3, generator=gen, device=self.device) * 8.0).to( + torch.float64 + ) + self.atype = torch.randint(0, 2, (1, n), generator=gen, device=self.device) + self.box = ( + (torch.eye(3, device=self.device) * 8.0).reshape(1, 9).to(torch.float64) + ) + + def _graph(self, des): + from deepmd.dpmodel.utils.neighbor_graph import ( + from_dense_quartet, + ) + + ec, ea, mp, nl = extend_input_and_build_neighbor_list( + self.coord, + self.atype, + des.get_rcut(), + des.get_sel(), + mixed_types=des.mixed_types(), + box=self.box, + ) + graph = from_dense_quartet(ec, nl, mp, compact=True) + return graph, self.atype.reshape(-1).to(self.device) + + def _assert_parity(self, des) -> None: + from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP + + self.assertTrue(des._fused_eligible("triton")) + tebd = des.type_embedding.call() + graph, atype_local = self._graph(des) + + def run(fn): + ev = graph.edge_vec.detach().clone().requires_grad_(True) + g = dataclasses.replace(graph, edge_vec=ev) + grrg, rot = fn(g) + (gvec,) = torch.autograd.grad(grrg.sum(), ev) + return grrg.detach(), rot.detach(), gvec.detach() + + d0, r0, g0 = run( + lambda g: DescrptDPA1DP.call_graph(des, g, atype_local, type_embedding=tebd) + ) + d1, r1, g1 = run( + lambda g: des._call_graph_triton(g, atype_local, type_embedding=tebd) + ) + torch.testing.assert_close(d1, d0, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(r1, r0, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(g1, g0, atol=1e-6, rtol=1e-6) + + def test_parity_concat_doubling(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [8, 16, 32])) + + def test_parity_concat_identity(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [8, 16, 16])) + + def test_parity_concat_one_side(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [8, 16, 32], one_side=True)) + + def test_parity_concat_silu(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [8, 16, 32], act="silu")) + + def test_parity_concat_nonpow2(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [12, 25, 50])) + + def _strip(self, neuron, **kw): + return _build_dpa1_expt(self.device, neuron, tebd_input_mode="strip", **kw) + + def test_parity_strip_doubling(self) -> None: + # Strip: the type feature enters as the multiplicative type-pair gate + # ``1 + tt[idx] * sw`` (smooth switch on by default); ``edge_conv`` runs + # with ``gated == 1``. + self._assert_parity(self._strip([8, 16, 32])) + + def test_parity_strip_identity(self) -> None: + self._assert_parity(self._strip([8, 16, 16])) + + def test_parity_strip_one_side(self) -> None: + # One-side folds the gate index by neighbor type alone (vs. the two-side + # (center, neighbor) pair). + self._assert_parity(self._strip([8, 16, 32], one_side=True)) + + def test_parity_strip_silu(self) -> None: + self._assert_parity(self._strip([8, 16, 32], act="silu")) + + def test_parity_strip_nonpow2(self) -> None: + self._assert_parity(self._strip([12, 25, 50])) + + def test_parity_strip_nonsmooth(self) -> None: + # Non-smooth strip drops the switch from the gate (``sw`` -> ones), the + # other ``gated == 1`` branch of ``_call_graph_triton``. + self._assert_parity(self._strip([8, 16, 32], smooth=False)) + + def test_make_fx_bakes_edge_conv(self) -> None: + des = _build_dpa1_expt(self.device, [8, 16, 32]) + tebd = des.type_embedding.call() + graph, atype_local = self._graph(des) + + def fn(edge_vec): + g = dataclasses.replace(graph, edge_vec=edge_vec) + return des._call_graph_triton(g, atype_local, type_embedding=tebd) + + traced = make_fx(fn, tracing_mode="real")(graph.edge_vec) + targets = [str(n.target) for n in traced.graph.nodes] + self.assertGreaterEqual(sum("edge_conv" in t for t in targets), 1) + + def test_make_fx_bakes_edge_conv_strip(self) -> None: + # Strip bakes both the gated ``edge_conv`` and the ``return_sw`` env-mat + # (the gate's per-edge switch) as opaque nodes in the graph .pt2. The + # env-mat routes to the opaque operator only at ``DP_TRITON_INFER >= 1`` + # (the export condition), so set the level while tracing. + des = self._strip([8, 16, 32]) + tebd = des.type_embedding.call() + graph, atype_local = self._graph(des) + + def fn(edge_vec): + g = dataclasses.replace(graph, edge_vec=edge_vec) + return des._call_graph_triton(g, atype_local, type_embedding=tebd) + + saved = os.environ.get("DP_TRITON_INFER") + try: + os.environ["DP_TRITON_INFER"] = "1" + traced = make_fx(fn, tracing_mode="real")(graph.edge_vec) + finally: + if saved is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved + targets = [str(n.target) for n in traced.graph.nodes] + self.assertGreaterEqual(sum("edge_conv" in t for t in targets), 1) + self.assertGreaterEqual(sum("edge_env_mat" in t for t in targets), 1) + + def test_call_graph_routes_on_level(self) -> None: + # call_graph bakes edge_conv only when DP_TRITON_INFER >= 1; level 0 + # keeps the dpmodel reference (no operator). + des = _build_dpa1_expt(self.device, [8, 16, 32]) + tebd = des.type_embedding.call() + graph, atype_local = self._graph(des) + saved = os.environ.get("DP_TRITON_INFER") + try: + counts = {} + for level in ("0", "1"): + os.environ["DP_TRITON_INFER"] = level + + def fn(edge_vec): + g = dataclasses.replace(graph, edge_vec=edge_vec) + return des.call_graph(g, atype_local, type_embedding=tebd) + + traced = make_fx(fn, tracing_mode="real")(graph.edge_vec) + counts[level] = sum( + "edge_conv" in str(n.target) for n in traced.graph.nodes + ) + finally: + if saved is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved + self.assertEqual(counts["0"], 0) + self.assertGreaterEqual(counts["1"], 1) + + def test_call_graph_fp32_edge_vec_cast(self) -> None: + # An fp32 model receives fp64 edge_vec at the graph .pt2 ABI. The public + # call_graph must cast it to the model precision (no ``double != float``) + # for both the reference and Triton routes, compute in fp32, and return + # the force gradient in the fp64 leaf dtype -- with the routes matching. + des = _build_dpa1_expt(self.device, [8, 16, 32], precision="float32") + tebd = des.type_embedding.call() + graph, atype_local = self._graph(des) + self.assertEqual(graph.edge_vec.dtype, torch.float64) + + def run(level): + os.environ["DP_TRITON_INFER"] = level + ev = graph.edge_vec.detach().clone().requires_grad_(True) + g = dataclasses.replace(graph, edge_vec=ev) + grrg, _ = des.call_graph(g, atype_local, type_embedding=tebd) + (gvec,) = torch.autograd.grad(grrg.sum(), ev) + return grrg.detach(), gvec.detach() + + saved = os.environ.get("DP_TRITON_INFER") + try: + d0, g0 = run("0") # dpmodel reference (edge_vec cast to fp32) + d1, g1 = run("1") # fused edge_conv (edge_vec cast to fp32) + finally: + if saved is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved + self.assertEqual(d1.dtype, torch.float32) # fp32 compute + self.assertEqual(g1.dtype, torch.float64) # force in the fp64 leaf dtype + self.assertTrue(torch.isfinite(g1).all()) + torch.testing.assert_close(d1, d0, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(g1, g0, atol=1e-5, rtol=1e-5) + + +@_GPU +class TestDpa1DenseRouting(unittest.TestCase): + """Dense-lower parity and make_fx bake of the fused env_mat + se_conv route. + + The dense ``call`` routes ``prod_env_mat`` through the fused + :func:`~deepmd.kernels.triton.env_mat.env_mat` operator (and the embedding + through :func:`~deepmd.kernels.triton.dpa1.se_conv.se_conv`) at + ``DP_TRITON_INFER >= 1``; this checks the operator wiring (parameters, atype + slicing) reproduces the dpmodel reference and bakes into the traced graph. + """ + + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + gen = torch.Generator(device=self.device).manual_seed(3) + n = 24 + self.coord = (torch.rand(1, n, 3, generator=gen, device=self.device) * 8.0).to( + torch.float64 + ) + self.atype = torch.randint(0, 2, (1, n), generator=gen, device=self.device) + self.box = ( + (torch.eye(3, device=self.device) * 8.0).reshape(1, 9).to(torch.float64) + ) + + def _dense_inputs(self, des): + ec, ea, mp, nl = extend_input_and_build_neighbor_list( + self.coord, + self.atype, + des.get_rcut(), + des.get_sel(), + mixed_types=des.mixed_types(), + box=self.box, + ) + return ec, ea, nl + + def _assert_parity(self, des) -> None: + self.assertTrue(des._fused_eligible("triton")) + ec, ea, nl = self._dense_inputs(des) + saved = os.environ.get("DP_TRITON_INFER") + + def run(level): + os.environ["DP_TRITON_INFER"] = level + c = ec.detach().clone().requires_grad_(True) + grrg = des.call(c, ea, nl)[0] + (gc,) = torch.autograd.grad(grrg.sum(), c) + return grrg.detach(), gc.detach() + + try: + d0, g0 = run("0") # dpmodel dense reference + d1, g1 = run("1") # fused env_mat + se_conv + finally: + if saved is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved + torch.testing.assert_close(d1, d0, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(g1, g0, atol=1e-6, rtol=1e-6) + + def test_parity_concat_doubling(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [8, 16, 32])) + + def test_parity_concat_identity(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [8, 16, 16])) + + def test_parity_concat_nonpow2(self) -> None: + self._assert_parity(_build_dpa1_expt(self.device, [12, 25, 50])) + + def test_make_fx_bakes_env_mat(self) -> None: + des = _build_dpa1_expt(self.device, [8, 16, 32]) + ec, ea, nl = self._dense_inputs(des) + saved = os.environ.get("DP_TRITON_INFER") + try: + os.environ["DP_TRITON_INFER"] = "1" + + def fn(c): + return des.call(c, ea, nl)[0] + + traced = make_fx(fn, tracing_mode="real")(ec) + finally: + if saved is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved + targets = [str(n.target) for n in traced.graph.nodes] + self.assertGreaterEqual(sum("env_mat.default" in t for t in targets), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/infer/test_graph_deepeval.py b/source/tests/pt_expt/infer/test_graph_deepeval.py index 07649d9e77..932fb25cfb 100644 --- a/source/tests/pt_expt/infer/test_graph_deepeval.py +++ b/source/tests/pt_expt/infer/test_graph_deepeval.py @@ -75,8 +75,8 @@ def _build_system( The blob keeps every atom within ``rcut`` of at most ``natoms - 1`` others (<< ``sel``), so the carry-all graph neighbor set equals the sel-capped - dense one. Varying ``natoms`` yields a different edge count, exercising the - DYNAMIC edge axis of the exported ``.pt2`` (B2.0). + dense one. Varying ``natoms`` yields a different edge count and exercises + the dynamic edge axis of the exported ``.pt2``. """ rng = np.random.default_rng(seed) box_size = 18.0 @@ -88,10 +88,7 @@ def _build_system( return coords, cells, atype -# Two DIFFERENT-size systems evaluated through the SAME exported ``.pt2``. -# Both are sparse, non-binding clusters but with different edge counts, so the -# second size FAILS against a static-``E`` artifact (B1) and PASSES only once -# the edge axis is dynamic (B2.0). +# Two system sizes with different edge counts use the same exported artifact. _SYSTEMS = { "small_8": {"natoms": 8, "seed": 20240626}, "large_20": {"natoms": 20, "seed": 20240701}, @@ -162,7 +159,7 @@ def graph_pt2(request): enumeration exports via unbacked SymInts (``xp_hint_dynamic_size``). ``smooth_type_embedding`` stays False: the smooth dense reference keeps sel-padding in its softmax denominator, so dense==carry-all parity holds - only for the non-smooth branch (PR-D divergence decision). + only for the non-smooth branch. The AOTI compile is slow (~90 s), so it is done once per param. The eager pt_expt model is returned alongside the archive path to serve as the dense @@ -197,9 +194,8 @@ def graph_pt2(request): def test_graph_pt2_deepeval_parity(graph_pt2, pbc, system) -> None: """Graph ``.pt2`` DeepEval == eager dense dpa1 (energy/force/virial), 1e-10. - Both ``_SYSTEMS`` are fed through the SAME module-scoped ``.pt2``; the - differing edge counts prove the exported artifact's edge axis is dynamic - (a static-``E`` B1 artifact would reject / mis-shape the larger system). + Both systems use the same module-scoped artifact; their different edge + counts exercise its dynamic edge axis. """ pt2_path, model = graph_pt2 coords, cells, atype = _build_system(**_SYSTEMS[system]) diff --git a/source/tests/pt_expt/model/test_dpa1_graph_lower.py b/source/tests/pt_expt/model/test_dpa1_graph_lower.py index d85e6dae27..bd77cc1783 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -113,8 +113,8 @@ def _make_model( attn_layer=attn_layer, attn_dotr=True, attn_mask=False, - # smooth attention keeps sel-padding in the dense softmax - # denominator; the carry-all graph drops it BY DESIGN (PR-D), so + # Smooth attention keeps sel-padding in the dense softmax + # denominator; the carry-all graph omits it, so # exact graph-vs-dense parity requires smooth=False here. smooth_type_embedding=smooth, activation_function="tanh", @@ -185,7 +185,7 @@ def _prepare_lower_inputs(self, periodic: bool): mapping_t = torch.tensor(mapping, dtype=torch.int64, device=self.device) return ext_coord, ext_atype, nlist_t, mapping_t - @pytest.mark.parametrize("attn_layer", [0, 2]) # factorizable AND attention + @pytest.mark.parametrize("attn_layer", [0, 2]) @pytest.mark.parametrize("periodic", [True, False]) # PBC vs non-PBC @pytest.mark.parametrize("do_av", [False, True]) # atom-virial off / on @pytest.mark.parametrize( @@ -239,6 +239,7 @@ def test_force_virial_parity_vs_legacy( graph = model.forward_common_lower_graph( atype_local, ng.n_node, + ng.n_node, ng.edge_index, ng.edge_vec, ng.edge_mask, @@ -303,23 +304,69 @@ def test_graph_lower_symbolic_trace(self, attn_layer) -> None: dtype=torch.float64, device=torch.device("cpu"), ) - atype, n_node, ei, ev, em, fp, ap, cs = sample + ( + atype, + n_node, + n_local, + ei, + ev, + em, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fp, + ap, + cs, + ) = sample traced = model.forward_lower_graph_exportable( atype, n_node, + n_local, ei, ev, em, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, fparam=fp, aparam=ap, do_atomic_virial=True, charge_spin=cs, + destination_sorted=True, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) - out = traced(atype, n_node, ei, ev, em, fp, ap, cs) + out = traced( + atype, + n_node, + n_local, + ei, + ev, + em, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fp, + ap, + cs, + ) ref = model.forward_common_lower_graph( - atype, n_node, ei, ev, em, fparam=fp, aparam=ap, do_atomic_virial=True + atype, + n_node, + n_local, + ei, + ev, + em, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam=fp, + aparam=ap, + do_atomic_virial=True, ) tol = {"rtol": 1e-12, "atol": 1e-12} torch.testing.assert_close(out["energy"], ref["energy_redu"], **tol) @@ -336,7 +383,7 @@ def test_smooth_attention_divergence_pinned(self) -> None: nonzero and bounded by the documented ~1e-4 magnitude. The carry-all graph drops sel-padding phantom terms from the smooth - attention softmax denominator BY DESIGN (NeighborGraph PR-D), while + attention softmax denominator, while the dense path keeps them, so dense output is sel-dependent. This test pins that divergence at the public model forward so a future refactor cannot silently change the carry-all smooth semantics. @@ -540,6 +587,8 @@ def test_graph_route_float32(self, attn_layer) -> None: ) tol = {"rtol": 1e-5, "atol": 1e-6} torch.testing.assert_close(graph["energy_redu"], dense["energy_redu"], **tol) + # The graph lower assembles force / virial in the model compute precision + # (fp32 here) while the dense reference stays fp64; compare values only. torch.testing.assert_close( - graph["energy_derv_r"], dense["energy_derv_r"], **tol + graph["energy_derv_r"], dense["energy_derv_r"], check_dtype=False, **tol ) diff --git a/source/tests/pt_expt/model/test_edge_energy_deriv.py b/source/tests/pt_expt/model/test_edge_energy_deriv.py index fafc8ac180..f03036633f 100644 --- a/source/tests/pt_expt/model/test_edge_energy_deriv.py +++ b/source/tests/pt_expt/model/test_edge_energy_deriv.py @@ -70,6 +70,49 @@ def test_padding_edges_contribute_nothing(self) -> None: torch.testing.assert_close(force, f2, rtol=1e-12, atol=1e-12) torch.testing.assert_close(gv, gv2, rtol=1e-12, atol=1e-12) + def test_force_precision_downcasts_inference_only(self) -> None: + """``force_precision`` assembles force / virial in that dtype for + inference, matching the fp64 assembly in content (the gradient carries + only the model precision); it is ignored when a graph is retained, so + training / double-backward keeps the fp64 leaf dtype. + """ + torch.manual_seed(2) + N = 6 + n_node = torch.tensor([N], dtype=torch.int64, device=env.DEVICE) + coord = torch.randn(N, 3, dtype=torch.float64, device=env.DEVICE) + src = torch.tensor([0, 1, 2, 3, 4, 5], device=env.DEVICE) + dst = torch.tensor([1, 0, 3, 2, 5, 4], device=env.DEVICE) + edge_index = torch.stack([src, dst], 0) + edge_mask = torch.ones(src.shape[0], dtype=torch.bool, device=env.DEVICE) + + def deriv(force_precision, create_graph): + ev = (coord[src] - coord[dst]).detach().requires_grad_(True) + energy = (torch.sin(ev).sum(-1) ** 2).sum() + return edge_energy_deriv( + energy, + ev, + edge_index, + edge_mask, + n_node, + do_atomic_virial=True, + create_graph=create_graph, + force_precision=force_precision, + ) + + f64, av64, gv64 = deriv(None, False) + f32, av32, gv32 = deriv(torch.float32, False) + # inference: force / virial are emitted in the requested precision ... + self.assertEqual(f32.dtype, torch.float32) + self.assertEqual(gv32.dtype, torch.float32) + self.assertEqual(av32.dtype, torch.float32) + # ... and match the fp64 assembly to the fp32 floor. + torch.testing.assert_close(f32.double(), f64, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(gv32.double(), gv64, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(av32.double(), av64, rtol=1e-5, atol=1e-5) + # a retained graph keeps the fp64 leaf dtype (downcast suppressed). + f_train, _, _ = deriv(torch.float32, True) + self.assertEqual(f_train.dtype, torch.float64) + def test_atom_virial_optional(self) -> None: """do_atomic_virial=False returns None for atom_virial; force+virial still computed.""" N = 3 diff --git a/source/tests/pt_expt/model/test_graph_builder_dispatch.py b/source/tests/pt_expt/model/test_graph_builder_dispatch.py index 6ec259c177..80d26a84fc 100644 --- a/source/tests/pt_expt/model/test_graph_builder_dispatch.py +++ b/source/tests/pt_expt/model/test_graph_builder_dispatch.py @@ -144,8 +144,7 @@ def test_dpmodel_backend_rejects_vesin(): def test_explicit_method_fails_fast_for_ineligible_descriptor(): """An EXPLICIT neighbor_graph_method must fail fast when the descriptor has no graph lower (mirrors the dpmodel guard; the default-path check in - _resolve_graph_method does not protect explicit methods). Regression for - OutisLi review on #5714. + _resolve_graph_method does not protect explicit methods). """ from deepmd.pt_expt.descriptor.se_e2_a import ( DescrptSeA, diff --git a/source/tests/pt_expt/model/test_graph_export.py b/source/tests/pt_expt/model/test_graph_export.py index 6b735aa3d5..1114b78848 100644 --- a/source/tests/pt_expt/model/test_graph_export.py +++ b/source/tests/pt_expt/model/test_graph_export.py @@ -60,54 +60,80 @@ def _graph_inputs(model): ) atype = torch.tensor([[0, 1, 0, 1, 0, 1]], dtype=torch.int64, device=env.DEVICE) box = torch.eye(3, dtype=torch.float64, device=env.DEVICE).reshape(1, 9) * 20.0 - g = build_neighbor_graph(coord, atype, box, model.get_rcut()) - return (atype.reshape(-1), g.n_node, g.edge_index, g.edge_vec, g.edge_mask) + g = build_neighbor_graph( + coord, + atype, + box, + model.get_rcut(), + canonicalize=True, + ) + edge_mask = g.edge_mask.clone() + assert bool(edge_mask[0]) + edge_mask[0] = False + return ( + atype.reshape(-1), + g.n_node, + g.n_node, + g.edge_index, + g.edge_vec, + edge_mask, + g.destination_order, + g.destination_row_ptr, + g.source_order, + g.source_row_ptr, + ) def test_graph_exportable_traces(): model = _model().eval() - atype, n_node, ei, ev, em = _graph_inputs(model) + graph_inputs = _graph_inputs(model) gm = model.forward_common_lower_graph_exportable( - atype, - n_node, - ei, - ev, - em, + *graph_inputs, do_atomic_virial=False, + destination_sorted=True, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) assert isinstance(gm, torch.nn.Module) # the traced module reproduces eager outputs - eager = model.forward_common_lower_graph( - atype, n_node, ei, ev, em, do_atomic_virial=False - ) - # traced module has placeholders for all 8 fn args (fparam/aparam/charge_spin=None) - traced = gm(atype, n_node, ei, ev, em, None, None, None) + eager = model.forward_common_lower_graph(*graph_inputs, do_atomic_virial=False) + # Optional conditioning inputs remain explicit placeholders. + traced = gm(*graph_inputs, None, None, None) # traced returns a tuple/dict; compare energy_redu te = traced["energy_redu"] if isinstance(traced, dict) else traced[1] torch.testing.assert_close(te, eager["energy_redu"], rtol=1e-10, atol=1e-10) +def test_graph_export_rejects_false_canonical_claim(): + model = _model().eval() + graph_inputs = list(_graph_inputs(model)) + graph_inputs[6] = torch.flip(graph_inputs[6], dims=(0,)) + + with pytest.raises(ValueError, match="destination_order"): + model.forward_common_lower_graph_exportable( + *graph_inputs, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + + @pytest.mark.parametrize("do_atomic_virial", [False, True]) # both branches of the bool def test_forward_lower_graph_exportable_public_keys(do_atomic_virial): """EnergyModel.forward_lower_graph_exportable: traces the public-key path and reproduces eager energy/force; atom_virial present iff do_atomic_virial. """ model = _model().eval() - atype, n_node, ei, ev, em = _graph_inputs(model) + graph_inputs = _graph_inputs(model) gm = model.forward_lower_graph_exportable( - atype, - n_node, - ei, - ev, - em, + *graph_inputs, do_atomic_virial=do_atomic_virial, + destination_sorted=True, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) assert isinstance(gm, torch.nn.Module) - out = gm(atype, n_node, ei, ev, em, None, None, None) + out = gm(*graph_inputs, None, None, None) # public key set (graph path is local-only: force/atom_virial, NOT extended_*) assert "atom_energy" in out and "energy" in out and "force" in out @@ -118,7 +144,7 @@ def test_forward_lower_graph_exportable_public_keys(do_atomic_virial): # values match the eager graph lower eager = model.forward_common_lower_graph( - atype, n_node, ei, ev, em, do_atomic_virial=do_atomic_virial + *graph_inputs, do_atomic_virial=do_atomic_virial ) torch.testing.assert_close( out["energy"], eager["energy_redu"], rtol=1e-10, atol=1e-10 diff --git a/source/tests/pt_expt/model/test_graph_ragged.py b/source/tests/pt_expt/model/test_graph_ragged.py index efe2ffeaec..ea066691fc 100644 --- a/source/tests/pt_expt/model/test_graph_ragged.py +++ b/source/tests/pt_expt/model/test_graph_ragged.py @@ -8,6 +8,9 @@ import torch +from deepmd.dpmodel.atomic_model.dp_atomic_model import ( + _extend_graph_aparam, +) from deepmd.pt.utils import ( env, ) @@ -99,6 +102,7 @@ def test_flat_energy_shapes(self) -> None: ret = self.model.forward_common_lower_graph( self.atype, self.n_node, + self.n_node, self.edge_index, self.edge_vec, self.edge_mask, @@ -128,6 +132,7 @@ def test_flat_atom_virial_shapes(self) -> None: ret = self.model.forward_common_lower_graph( self.atype, self.n_node, + self.n_node, self.edge_index, self.edge_vec, self.edge_mask, @@ -143,16 +148,86 @@ def test_flat_atom_virial_shapes(self) -> None: assert torch.isfinite(ret["energy_derv_c"]).all() assert torch.isfinite(ret["energy_derv_c_redu"]).all() + def test_halo_nodes_do_not_contribute_atomic_energy(self) -> None: + n_local = torch.tensor([2, 1], dtype=torch.int64, device=self.device) + ret = self.model.forward_common_lower_graph( + self.atype, + self.n_node, + n_local, + self.edge_index, + self.edge_vec, + self.edge_mask, + do_atomic_virial=True, + ) + + torch.testing.assert_close( + ret["energy"][[2, 4]], + torch.zeros_like(ret["energy"][[2, 4]]), + ) + torch.testing.assert_close( + ret["energy_redu"][0], + ret["energy"][:2].sum(dim=0), + ) + torch.testing.assert_close( + ret["energy_redu"][1], + ret["energy"][3:4].sum(dim=0), + ) + torch.testing.assert_close( + ret["mask"], + torch.tensor([1, 1, 0, 1, 0], dtype=torch.int32, device=self.device), + ) + + def test_ragged_local_aparam_expands_to_halo_axis(self) -> None: + n_node = torch.tensor([3, 4], dtype=torch.int64, device=self.device) + n_local = torch.tensor([2, 1], dtype=torch.int64, device=self.device) + aparam = torch.tensor( + [ + [[1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [99.0, 99.0]], + ], + dtype=torch.float64, + device=self.device, + ) + actual = _extend_graph_aparam(aparam, n_node, n_local, 7) + expected = torch.tensor( + [ + [1.0, 2.0], + [3.0, 4.0], + [0.0, 0.0], + [5.0, 6.0], + [0.0, 0.0], + [0.0, 0.0], + [0.0, 0.0], + ], + dtype=torch.float64, + device=self.device, + ) + torch.testing.assert_close(actual, expected) + + def test_zero_owned_aparam_expands_to_zeros(self) -> None: + n_node = torch.tensor([3], dtype=torch.int64, device=self.device) + n_local = torch.tensor([0], dtype=torch.int64, device=self.device) + aparam = torch.empty( + (1, 0, 2), + dtype=torch.float64, + device=self.device, + ) + actual = _extend_graph_aparam(aparam, n_node, n_local, 3) + torch.testing.assert_close( + actual, + torch.zeros((3, 2), dtype=torch.float64, device=self.device), + ) + def test_invariant_to_charge_spin(self) -> None: """dpa1 does NOT consume charge_spin (``get_dim_chg_spin() == 0``); - forward_common_lower_graph accepts it only for ABI stability with - charge/spin descriptors (dpa3/dpa4, PR-G), so energy / force / virial / - atom-virial must be INVARIANT to it. + forward_common_lower_graph accepts it for descriptors with charge/spin + conditioning, so dpa1 outputs must be invariant to it. """ assert self.model.get_descriptor().get_dim_chg_spin() == 0 # dpa1 args = ( self.atype, self.n_node, + self.n_node, self.edge_index, self.edge_vec, self.edge_mask, diff --git a/source/tests/pt_expt/utils/test_canonical_graph.py b/source/tests/pt_expt/utils/test_canonical_graph.py new file mode 100644 index 0000000000..e7041d7f6e --- /dev/null +++ b/source/tests/pt_expt/utils/test_canonical_graph.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import pytest +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, +) +from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + validate_canonical_graph_shapes, +) + + +def _generic_graph(physical_edge_count: int) -> NeighborGraph: + storage_count = physical_edge_count + 2 + edge_index = torch.zeros((2, storage_count), dtype=torch.int64) + edge_vec = torch.zeros((storage_count, 3), dtype=torch.float64) + edge_mask = torch.zeros(storage_count, dtype=torch.bool) + if physical_edge_count: + edge_mask[:physical_edge_count] = True + edge_index[0, :physical_edge_count] = 0 + edge_index[1, :physical_edge_count] = 0 + edge_vec[:physical_edge_count, 0] = 1.5 + row_ptr = torch.tensor([0, physical_edge_count], dtype=torch.int64) + return NeighborGraph( + n_node=torch.tensor([1], dtype=torch.int64), + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + n_local=torch.tensor([1], dtype=torch.int64), + destination_order=torch.arange(storage_count, dtype=torch.int64), + destination_row_ptr=row_ptr, + source_order=torch.arange(storage_count, dtype=torch.int64), + source_row_ptr=row_ptr.clone(), + destination_sorted=True, + ) + + +@pytest.mark.parametrize("physical_edge_count", [0, 1]) +def test_storage_guards_remain_outside_csr(physical_edge_count: int) -> None: + compact = canonical_graph_from_neighbor_graph(_generic_graph(physical_edge_count)) + 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 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) diff --git a/source/tests/pt_expt/utils/test_edge_env_mat_triton.py b/source/tests/pt_expt/utils/test_edge_env_mat_triton.py new file mode 100644 index 0000000000..62cedf294c --- /dev/null +++ b/source/tests/pt_expt/utils/test_edge_env_mat_triton.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Graph-native (edge-stream) environment-matrix Triton kernel. + +The edge form (:func:`deepmd.kernels.triton.env_mat.edge_env_mat`) is the +slot-free analogue used only by the pt_expt graph lower: the relative vector +``edge_vec`` is given directly (no neighbor gather) and the backward +differentiates ``edge_vec`` (the graph-path force leaf), so no scatter is +needed. These tests check parity against the dpmodel reference +(:func:`deepmd.dpmodel.utils.neighbor_graph.env.edge_env_mat`) in fp32 and fp64, +including padding (zero-vector) edges, the optionally returned smooth switch +``sw`` and its gradient (the strip type-pair gate consumes it), and +composability under ``make_fx``. +""" + +import os +import unittest + +import torch + +from deepmd.kernels.triton.env_mat import ( + TRITON_AVAILABLE, + edge_env_mat, +) + +_GPU = unittest.skipUnless( + torch.cuda.is_available() and TRITON_AVAILABLE, + "CUDA and Triton are required for the edge env_mat kernel", +) + + +def _make_edges(dtype, device, n_edge=4000, ntypes=3, seed=0): + """Random edge stream with padding (zero-vector) edges, as on the graph path.""" + g = torch.Generator(device=device).manual_seed(seed) + edge_vec = torch.randn(n_edge, 3, generator=g, device=device, dtype=dtype) * 2.0 + edge_mask = torch.rand(n_edge, generator=g, device=device) > 0.2 + edge_vec[~edge_mask] = 0.0 + center_type = torch.randint(0, ntypes, (n_edge,), generator=g, device=device) + davg = torch.randn(ntypes, 4, generator=g, device=device, dtype=dtype) + dstd = 0.5 + torch.rand(ntypes, 4, generator=g, device=device, dtype=dtype) + return edge_vec, center_type, edge_mask, davg, dstd + + +@_GPU +class TestEdgeEnvMatTriton(unittest.TestCase): + def setUp(self) -> None: + torch.backends.cuda.matmul.allow_tf32 = False + self.device = torch.device("cuda") + + def tearDown(self) -> None: + os.environ.pop("DP_TRITON_INFER", None) + + def test_forward_and_edge_grad_parity(self) -> None: + from deepmd.dpmodel.utils.neighbor_graph.env import edge_env_mat as edge_ref + + rcut, rcut_smth = 6.0, 2.0 + for dtype in (torch.float32, torch.float64): + with self.subTest(dtype=dtype): + ev, ct, mask, davg, dstd = _make_edges(dtype, self.device) + gout = torch.randn(ev.shape[0], 4, device=self.device, dtype=dtype) + + os.environ["DP_TRITON_INFER"] = "1" + e1 = ev.clone().requires_grad_() + r1 = edge_env_mat(e1, ct, davg, dstd, rcut, rcut_smth, edge_mask=mask) + (g1,) = torch.autograd.grad((r1 * gout).sum(), e1) + + e0 = ev.clone().requires_grad_() + r0 = edge_ref(e0, ct, davg, dstd, rcut, rcut_smth, edge_mask=mask) + (g0,) = torch.autograd.grad((r0 * gout).sum(), e0) + + ftol = 1e-5 if dtype == torch.float32 else 1e-10 + gtol = 1e-3 if dtype == torch.float32 else 1e-8 + self.assertLess((r1 - r0).abs().max().item(), ftol) + self.assertLess( + (g1 - g0).abs().max().item() / (g0.abs().max().item() + 1e-30), + gtol, + ) + + def test_return_sw_parity(self) -> None: + # ``return_sw`` also emits the per-edge switch (strip gate input). Feed + # independent cotangents into ``env`` and ``sw`` so the backward + # exercises the env path, the switch path (the fused ``g_sw * s'`` term) + # and their sum into the ``edge_vec`` leaf. + from deepmd.dpmodel.utils.neighbor_graph.env import edge_env_mat as edge_ref + + rcut, rcut_smth = 6.0, 2.0 + for dtype in (torch.float32, torch.float64): + with self.subTest(dtype=dtype): + ev, ct, mask, davg, dstd = _make_edges(dtype, self.device) + g_env = torch.randn(ev.shape[0], 4, device=self.device, dtype=dtype) + g_sw = torch.randn(ev.shape[0], 1, device=self.device, dtype=dtype) + + os.environ["DP_TRITON_INFER"] = "1" + e1 = ev.clone().requires_grad_() + env1, sw1 = edge_env_mat( + e1, ct, davg, dstd, rcut, rcut_smth, edge_mask=mask, return_sw=True + ) + loss1 = (env1 * g_env).sum() + (sw1 * g_sw).sum() + (g1,) = torch.autograd.grad(loss1, e1) + + e0 = ev.clone().requires_grad_() + env0, sw0 = edge_ref( + e0, ct, davg, dstd, rcut, rcut_smth, edge_mask=mask, return_sw=True + ) + loss0 = (env0 * g_env).sum() + (sw0 * g_sw).sum() + (g0,) = torch.autograd.grad(loss0, e0) + + ftol = 1e-5 if dtype == torch.float32 else 1e-10 + gtol = 1e-3 if dtype == torch.float32 else 1e-8 + self.assertEqual(sw1.shape, (ev.shape[0], 1)) + self.assertLess((sw1 - sw0).abs().max().item(), ftol) + self.assertLess((env1 - env0).abs().max().item(), ftol) + self.assertLess( + (g1 - g0).abs().max().item() / (g0.abs().max().item() + 1e-30), + gtol, + ) + + def test_make_fx_opaque_operator(self) -> None: + from torch.fx.experimental.proxy_tensor import ( + make_fx, + ) + + os.environ["DP_TRITON_INFER"] = "1" + ev, ct, mask, davg, dstd = _make_edges(torch.float32, self.device) + + def fn(edge_vec, center_type, edge_mask, avg, std): + return edge_env_mat( + edge_vec, center_type, avg, std, 6.0, 2.0, edge_mask=edge_mask + ).sum() + + gm = make_fx(fn, tracing_mode="fake")(ev, ct, mask, davg, dstd) + targets = [str(n.target) for n in gm.graph.nodes] + self.assertTrue(any("edge_env_mat.default" in t for t in targets)) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/utils/test_graph_csr.py b/source/tests/pt_expt/utils/test_graph_csr.py new file mode 100644 index 0000000000..be3eb66f13 --- /dev/null +++ b/source/tests/pt_expt/utils/test_graph_csr.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for PyTorch graph-lower CSR export validation.""" + +import pytest +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + build_edge_csr, +) +from deepmd.pt_expt.utils.graph_csr import ( + validate_graph_csr_for_export, +) + + +def _csr_inputs(canonicalize: bool) -> tuple[torch.Tensor, ...]: + edge_index = torch.tensor( + [[1, 2, 0, 0], [0, 0, 1, 0]], + dtype=torch.int64, + ) + edge_vec = torch.zeros(4, 3, dtype=torch.float64) + edge_mask = torch.tensor([True, True, True, False]) + ( + edge_index, + _edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = build_edge_csr( + edge_index, + edge_vec, + edge_mask, + n_nodes=3, + canonicalize=canonicalize, + ) + return ( + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) + + +def test_accepts_canonical_csr() -> None: + inputs = _csr_inputs(canonicalize=True) + + validate_graph_csr_for_export( + *inputs, + node_count=3, + destination_sorted=True, + ) + + +def test_accepts_permutation_csr() -> None: + inputs = _csr_inputs(canonicalize=False) + + validate_graph_csr_for_export( + *inputs, + node_count=3, + destination_sorted=False, + ) + + +def test_rejects_nonidentity_canonical_order() -> None: + ( + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = _csr_inputs(canonicalize=True) + destination_order = destination_order.clone() + destination_order[:2] = torch.tensor([1, 0], dtype=torch.int64) + + with pytest.raises(ValueError, match="identity destination_order"): + validate_graph_csr_for_export( + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + node_count=3, + destination_sorted=True, + ) + + +def test_rejects_active_edge_outside_its_csr_row() -> None: + ( + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + ) = _csr_inputs(canonicalize=True) + destination_row_ptr = destination_row_ptr.clone() + destination_row_ptr[1] = 1 + + with pytest.raises(ValueError, match="destination_order entries outside"): + validate_graph_csr_for_export( + edge_index, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + node_count=3, + destination_sorted=True, + ) diff --git a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py index 54aa9f688d..cadec99c89 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -15,8 +15,11 @@ import zipfile import pytest +import torch from deepmd.pt_expt.utils.serialization import ( + _graph_edge_dtype, + _supports_graph_export, deserialize_to_file, ) @@ -89,11 +92,42 @@ def test_graph_pt2_has_lower_input_kind_graph(dpa1_dpmodel_data) -> None: ) meta = _read_metadata(p) assert meta["lower_input_kind"] == "graph" - # B2.0: the edge axis is DYNAMIC (Dim("nedge", min=2)); there is no static - # capacity baked into the AOTI artifact, so no ``edge_capacity`` is persisted. + assert meta["graph_edge_dtype"] == "float64" + # A dynamic edge axis has no persisted static capacity. assert "edge_capacity" not in meta +@pytest.mark.parametrize( + ("statistics_dtype", "expected"), + [(torch.float32, "float32"), (torch.float64, "float64")], +) +def test_compressed_graph_uses_compute_precision_edge_geometry( + statistics_dtype: torch.dtype, expected: str +) -> None: + """Compressed DPA1 graph geometry follows descriptor compute precision.""" + + class _Descriptor: + geo_compress = True + + class _Block: + mean = torch.empty(0, dtype=statistics_dtype) + + se_atten = _Block() + + def _fused_eligible(self, backend: str) -> bool: + return backend == "cuda" and self.se_atten.mean.dtype == torch.float32 + + class _AtomicModel: + descriptor = _Descriptor() + + class _Model: + atomic_model = _AtomicModel() + + assert _graph_edge_dtype(_Model(), "graph") == expected + assert _graph_edge_dtype(_Model(), "nlist") == "float64" + assert _supports_graph_export(_Model()) is (statistics_dtype == torch.float32) + + def test_dense_pt2_has_lower_input_kind_nlist(dpa1_dpmodel_data) -> None: """Default (``lower_kind="nlist"``) -> metadata ``lower_input_kind == "nlist"``.""" with tempfile.TemporaryDirectory() as d: @@ -105,6 +139,7 @@ def test_dense_pt2_has_lower_input_kind_nlist(dpa1_dpmodel_data) -> None: ) meta = _read_metadata(p) assert meta["lower_input_kind"] == "nlist" + assert meta["graph_edge_dtype"] == "float64" # edge_capacity is a graph-only artifact constant; the dense path omits it. assert "edge_capacity" not in meta @@ -114,8 +149,7 @@ def test_neighbor_graph_method_rejected_on_nlist_artifact(dpa1_dpmodel_data) -> The knob is consumed only by graph-form ``.pt2`` eval; silently ignoring it on nlist-form artifacts misled users into thinking they selected an - O(N) builder (OutisLi review, #5714). The nlist-path knob is - ``nlist_backend``. + O(N) builder. The nlist-path knob is ``nlist_backend``. """ from deepmd.infer import ( DeepPot, @@ -130,7 +164,7 @@ def test_neighbor_graph_method_rejected_on_nlist_artifact(dpa1_dpmodel_data) -> ) with pytest.raises(ValueError, match="graph-form"): DeepPot(p, neighbor_graph_method="vesin") - # the default stays accepted (no behavior change) + # The default remains valid for nlist-form artifacts. DeepPot(p) diff --git a/source/tests/test_convert_backend.py b/source/tests/test_convert_backend.py new file mode 100644 index 0000000000..063caec868 --- /dev/null +++ b/source/tests/test_convert_backend.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import pytest + +from deepmd.backend.backend import ( + Backend, +) +from deepmd.entrypoints.convert_backend import ( + convert_backend, +) + + +def test_convert_backend_automatically_selects_lower_kind( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + class InputBackend: + name = "input" + + @staticmethod + def serialize_hook(path: str) -> dict[str, str]: + return {"path": path} + + class OutputBackend: + name = "output" + + @staticmethod + def deserialize_hook( + path: str, + data: dict[str, str], + *, + lower_kind: str = "nlist", + do_atomic_virial: bool = False, + ) -> None: + captured.update( + path=path, + data=data, + lower_kind=lower_kind, + do_atomic_virial=do_atomic_virial, + ) + + def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: + return InputBackend if path.endswith(".input") else OutputBackend + + monkeypatch.setattr(Backend, "detect_backend_by_model", detect_backend) + + convert_backend(INPUT="model.input", OUTPUT="model.output") + + assert captured["lower_kind"] == "auto" + assert captured["do_atomic_virial"] is False