From bcf3c54043a3e44b126a7279dbcb9001a16af180 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 4 Jul 2026 19:35:59 +0800 Subject: [PATCH 1/9] perf(pt/dpa1): optimize graph CUDA inference and deployment DPA1 graph inference previously decomposed descriptor, fitting, and force/virial work into many generic tensor operations, retained large edge-scale autograd state, and lacked a shared sparse topology contract across Python export, the C++ runtime, and LAMMPS. Build an inference-oriented graph pipeline that keeps eligible attention-free DPA1 models and their neighbor data on the GPU while retaining explicit fallbacks for unsupported configurations. - Add inference-only CUDA operators for attention-free concat and strip DPA1. The uncompressed path fuses environment construction, the three-layer embedding network, type handling, moment reduction, Gram contraction, and analytical edge-gradient backward. The geometrically compressed strip path fuses quintic table interpolation, type-pair gating, moments, descriptors, and analytical gradients. It supports masked edges, canonical or permutation CSR, int32 or int64 addressing, width buckets from 8 through 256, non-bucket padding and slicing, and axis widths up to 16. - Add descriptor-agnostic fused energy fitting and force/virial operators. Run eligible fitting networks through pedantic FP32 cuBLAS GEMMs with TF32 disabled and fused bias, activation, timestep, residual, and derivative epilogues, while retaining atomic and per-frame energies in FP64. Reduce both edge incidences with one warp per node and no global floating-point atomics, then accumulate frame virials through FP64 partial reductions without materializing an `(E, 9)` outer-product tensor. - Define cumulative `DP_CUDA_INFER` levels. Level 1 uses separately registered descriptor, fitting, and CSR force operators with first-order autograd. Level 2 returns force and virial as values, using one opaque end-to-end operator for eligible uncompressed DPA1 and an explicit custom-operator chain for compressed DPA1. This removes the inference autograd tape, suppresses the unused rotation output, avoids retaining the descriptor only for shape metadata, and shortens saved-state lifetimes for memory reuse. - Add `make_fx`- and export-composable Triton kernels for the PyTorch `se`-family environment matrix and the dense and graph DPA1 environment convolutions. Their closed-form first backward preserves the force path while avoiding decomposed gather, switch, normalization, gating, and segment reduction chains. `DP_TRITON_INFER=2` resolves per-GPU launch tables and freeze-time tuning fills uncovered model shapes; level 3 optionally enables the compensated FP16x3 final embedding GEMM. The default CUDA graph path remains FP32. - Select compressed-kernel resource policy at first uncaptured use from balanced or occupancy-oriented 128/256-thread launches, cache the result by device, direction, descriptor shape, topology, index width, and workload class, and retain architecture-specific fallbacks for small inputs, capture, or tuning failures. Build the native operators only with a CUDA-enabled PyTorch and include a lowest-supported portable PTX target alongside native code. - Extend `NeighborGraph` with keyword-only destination/source CSR orders, int64 row pointers, and an explicit `destination_sorted` property without changing its original positional constructor. Generic builders preserve payload order and omit CSR by default; consumers opt into CSR or stable destination-major canonicalization, which moves masked guards to the suffix and makes destination order the identity. Edge masks remain authoritative inside every row, and export validates permutations, row membership, bounds, and canonical identity before tracing. - Make graph-form AOTInductor deployment select the optimized path automatically for eligible `.pt2` conversion and export compressed graph models directly. Keep frame, node, and edge axes dynamic; preserve custom operators through export; record the graph edge-vector dtype in metadata; and make per-atom virial part of the graph artifact contract required by the Kokkos consumer. Eligible compressed artifacts accept FP32 edge geometry directly, while generic and uncompressed graph artifacts retain the FP64 geometry ABI. Use int64 Inductor indexing and bounded Triton launch tiling for large dynamic graphs. - Extend DeepEval and the C++ API to consume the same canonical graph ABI. Add device-edge capability and dtype queries, FP32 and FP64 edge-vector overloads, runtime frame and atomic parameters, total-versus-owned node counts, and optional communication metadata. Host ingestion canonicalizes arbitrary payloads, while the device path constructs destination identity CSR and source order with histogram, prefix-sum, and counting scatter. - Add `pair_style deepmd/kk` for device-resident edge and graph `.pt2` inference. Build compact model-cutoff edges from the Kokkos full neighbor list with count, scan, and fill passes; compact NULL-mapped atom types; emit FP32 or FP64 vectors according to artifact metadata; and scatter model-node outputs back on device. Single-rank execution folds periodic images onto local owners, whereas domain decomposition retains local-plus-halo nodes and explicitly reverse-communicates force and centroid per-atom virial through either device-aware or host-staged communication. Message-passing edge artifacts use their with-comm forward, including empty-rank phantom inputs; message-passing graph artifacts fail fast because that multi-rank ABI is not supported. - Preserve graph correctness at ownership and masking boundaries. Exclude halo fitting outputs from energy, zero halo atomic parameters, retain virtual-atom and pair-exclusion masks in fused level 2, handle zero-node graphs, and reduce frame virial from node virials instead of a contended edge scatter. Validate and broadcast multi-frame `fparam`/`aparam` layouts, fix wide atomic-parameter allocation when filtering virtual atoms, and avoid copying incompletely constructed LAMMPS model members. Reduce graph-builder overhead with bounded GPU neighbor-capacity estimates and direct length-three periodic-shift reductions. - Require regeneration of graph-form `.pt2` artifacts because their positional ABI now includes `n_local` and both CSR views. Compressed graph export is limited to FP32, geometrically compressed, attention-free strip models without excluded type pairs, with output width at most 256 and `axis_neuron <= 16`. The uncompressed CUDA descriptor requires three FP32 embedding layers whose widths stay equal or double within its compiled bounds. Other configurations continue through Triton, reference graph, or nlist lowering. `deepmd/kk` additionally requires a GPU edge/graph artifact, one model, and a valid atom map. - Add regression coverage for CUDA and Triton forward/gradient parity, compressed and uncompressed end-to-end energy, force, global virial, atomic energy, and atom virial; smooth and non-smooth one- and two-sided type gates; residual and activation variants; non-power-of-two widths; int32/int64 addressing; masked cached edges; canonical and permuted CSR; ownership; zero-node and many-frame reductions; `make_fx`, `torch.compile`, and dynamic graph export. Extend DeepEval and C++ tests for dynamic edge counts, multi-rank graph ingestion, parameter validation and broadcasting, atomic outputs, CSR construction, and virtual-atom parameter preservation. --- .../dpmodel/atomic_model/base_atomic_model.py | 61 +- .../dpmodel/atomic_model/dp_atomic_model.py | 61 +- deepmd/dpmodel/descriptor/dpa1.py | 25 +- deepmd/dpmodel/model/make_model.py | 27 +- .../dpmodel/utils/neighbor_graph/__init__.py | 13 +- deepmd/dpmodel/utils/neighbor_graph/angles.py | 10 +- .../utils/neighbor_graph/ase_builder.py | 20 +- .../dpmodel/utils/neighbor_graph/builder.py | 104 + deepmd/dpmodel/utils/neighbor_graph/csr.py | 190 ++ .../utils/neighbor_graph/derivatives.py | 67 +- .../dpmodel/utils/neighbor_graph/from_ijs.py | 40 + deepmd/dpmodel/utils/neighbor_graph/graph.py | 47 +- deepmd/dpmodel/utils/neighbor_graph/pairs.py | 8 +- deepmd/entrypoints/convert_backend.py | 28 +- deepmd/kernels/autotune.py | 87 + deepmd/kernels/cuda/__init__.py | 24 + deepmd/kernels/cuda/dpa1/__init__.py | 10 + deepmd/kernels/cuda/dpa1/graph_compress.py | 995 ++++++++++ deepmd/kernels/cuda/dpa1/graph_descriptor.py | 763 ++++++++ .../kernels/cuda/dpa1/graph_energy_force.py | 433 +++++ deepmd/kernels/cuda/edge_force_virial.py | 178 ++ deepmd/kernels/cuda/graph_fitting.py | 312 +++ deepmd/kernels/triton/dpa1/__init__.py | 25 + deepmd/kernels/triton/dpa1/activation.py | 61 + deepmd/kernels/triton/dpa1/edge_conv.py | 900 +++++++++ deepmd/kernels/triton/dpa1/gemm_fp16x3.py | 267 +++ deepmd/kernels/triton/dpa1/se_conv.py | 849 ++++++++ .../kernels/triton/dpa1/sweep_tile_configs.py | 306 +++ deepmd/kernels/triton/dpa1/tile_configs.py | 166 ++ deepmd/kernels/triton/env_mat.py | 1116 +++++++++++ deepmd/kernels/utils.py | 115 +- deepmd/pt/model/descriptor/env_mat.py | 27 + deepmd/pt/model/descriptor/se_atten.py | 200 +- deepmd/pt_expt/descriptor/dpa1.py | 957 ++++++++- deepmd/pt_expt/descriptor/se_atten_v2.py | 59 + deepmd/pt_expt/entrypoints/compress.py | 50 +- deepmd/pt_expt/fitting/ener_fitting.py | 49 + deepmd/pt_expt/infer/deep_eval.py | 99 +- deepmd/pt_expt/model/edge_transform_output.py | 74 +- deepmd/pt_expt/model/ener_model.py | 65 +- deepmd/pt_expt/model/make_model.py | 217 ++- deepmd/pt_expt/train/training.py | 25 + deepmd/pt_expt/utils/edge_schema.py | 7 +- deepmd/pt_expt/utils/graph_csr.py | 173 ++ deepmd/pt_expt/utils/nv_graph_builder.py | 40 +- deepmd/pt_expt/utils/serialization.py | 342 +++- deepmd/pt_expt/utils/vesin_graph_builder.py | 48 +- deepmd/pt_expt/utils/vesin_neighbor_list.py | 11 +- pyproject.toml | 2 +- source/CMakeLists.txt | 10 + source/api_cc/include/DeepPot.h | 101 +- source/api_cc/include/DeepPotPTExpt.h | 135 +- source/api_cc/include/commonPT.h | 100 +- source/api_cc/src/DeepPot.cc | 102 +- source/api_cc/src/DeepPotPTExpt.cc | 720 +++++-- source/api_cc/src/common.cc | 2 +- source/api_cc/src/commonPTExpt.h | 11 +- ..._deeppot_a_fparam_aparam_nframes_ptexpt.cc | 45 + .../tests/test_deeppot_dpa1_graph_ptexpt.cc | 103 +- source/api_cc/tests/test_deeppot_ptexpt.cc | 2 +- .../api_cc/tests/test_neighbor_list_data.cc | 50 +- source/api_cc/tests/test_select_real_atoms.cc | 40 + source/lmp/builtin.cmake | 18 +- source/lmp/pair_base.cpp | 14 +- source/lmp/pair_base.h | 10 +- source/lmp/pair_deepmd.cpp | 45 +- source/lmp/pair_deepmd.h | 4 + source/lmp/pair_deepmd_kokkos.cpp | 767 ++++++++ source/lmp/pair_deepmd_kokkos.h | 137 ++ source/lmp/pair_deepspin.cpp | 33 +- source/op/pt/CMakeLists.txt | 59 +- source/op/pt/dpa1_graph_common.cuh | 564 ++++++ source/op/pt/dpa1_graph_compress.cu | 1512 +++++++++++++++ source/op/pt/dpa1_graph_compress_tuning.h | 229 +++ source/op/pt/dpa1_graph_descriptor.cu | 1499 ++++++++++++++ source/op/pt/dpa1_graph_energy_force.cu | 156 ++ source/op/pt/edge_force_virial.cu | 405 ++++ source/op/pt/graph_fitting.cu | 358 ++++ source/op/pt/graph_ops.h | 130 ++ .../common/dpmodel/test_call_lower_graph.py | 4 +- .../dpmodel/test_dpa1_call_graph_block.py | 3 - .../test_dpa1_call_graph_descriptor.py | 4 +- .../test_dpa1_graph_attention_parity.py | 11 +- .../dpmodel/test_dpa1_graph_model_energy.py | 4 +- .../common/dpmodel/test_edge_force_virial.py | 35 +- .../common/dpmodel/test_neighbor_graph.py | 119 ++ .../dpmodel/test_neighbor_graph_builder.py | 39 + .../pt/model/test_descriptor_dpa1_triton.py | 593 ++++++ source/tests/pt/model/test_env_mat_triton.py | 189 ++ .../pt_expt/descriptor/test_dpa1_cuda.py | 1721 +++++++++++++++++ .../pt_expt/descriptor/test_dpa1_triton.py | 370 ++++ .../pt_expt/infer/test_graph_deepeval.py | 16 +- .../pt_expt/model/test_dpa1_graph_lower.py | 67 +- .../pt_expt/model/test_edge_energy_deriv.py | 43 + .../model/test_graph_builder_dispatch.py | 3 +- .../tests/pt_expt/model/test_graph_export.py | 68 +- .../tests/pt_expt/model/test_graph_ragged.py | 81 +- .../pt_expt/utils/test_edge_env_mat_triton.py | 136 ++ source/tests/pt_expt/utils/test_graph_csr.py | 114 ++ .../pt_expt/utils/test_graph_pt2_metadata.py | 44 +- 100 files changed, 19890 insertions(+), 888 deletions(-) create mode 100644 deepmd/dpmodel/utils/neighbor_graph/csr.py create mode 100644 deepmd/kernels/autotune.py create mode 100644 deepmd/kernels/cuda/__init__.py create mode 100644 deepmd/kernels/cuda/dpa1/__init__.py create mode 100644 deepmd/kernels/cuda/dpa1/graph_compress.py create mode 100644 deepmd/kernels/cuda/dpa1/graph_descriptor.py create mode 100644 deepmd/kernels/cuda/dpa1/graph_energy_force.py create mode 100644 deepmd/kernels/cuda/edge_force_virial.py create mode 100644 deepmd/kernels/cuda/graph_fitting.py create mode 100644 deepmd/kernels/triton/dpa1/__init__.py create mode 100644 deepmd/kernels/triton/dpa1/activation.py create mode 100644 deepmd/kernels/triton/dpa1/edge_conv.py create mode 100644 deepmd/kernels/triton/dpa1/gemm_fp16x3.py create mode 100644 deepmd/kernels/triton/dpa1/se_conv.py create mode 100644 deepmd/kernels/triton/dpa1/sweep_tile_configs.py create mode 100644 deepmd/kernels/triton/dpa1/tile_configs.py create mode 100644 deepmd/kernels/triton/env_mat.py create mode 100644 deepmd/pt_expt/utils/graph_csr.py create mode 100644 source/api_cc/tests/test_select_real_atoms.cc create mode 100644 source/lmp/pair_deepmd_kokkos.cpp create mode 100644 source/lmp/pair_deepmd_kokkos.h create mode 100644 source/op/pt/dpa1_graph_common.cuh create mode 100644 source/op/pt/dpa1_graph_compress.cu create mode 100644 source/op/pt/dpa1_graph_compress_tuning.h create mode 100644 source/op/pt/dpa1_graph_descriptor.cu create mode 100644 source/op/pt/dpa1_graph_energy_force.cu create mode 100644 source/op/pt/edge_force_virial.cu create mode 100644 source/op/pt/graph_fitting.cu create mode 100644 source/op/pt/graph_ops.h create mode 100644 source/tests/pt/model/test_descriptor_dpa1_triton.py create mode 100644 source/tests/pt/model/test_env_mat_triton.py create mode 100644 source/tests/pt_expt/descriptor/test_dpa1_cuda.py create mode 100644 source/tests/pt_expt/descriptor/test_dpa1_triton.py create mode 100644 source/tests/pt_expt/utils/test_edge_env_mat_triton.py create mode 100644 source/tests/pt_expt/utils/test_graph_csr.py diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index bf41735f89..f72ff27380 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -313,6 +313,42 @@ 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_inputs( + self, + graph: "NeighborGraph", + atype: Array, + ) -> tuple["NeighborGraph", Array, Array]: + """Apply graph masks shared by standard and fused atomic forwards.""" + 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)) + output_mask = atom_mask + if graph.n_local is not None: + from deepmd.dpmodel.utils.neighbor_graph import ( + node_ownership_mask, + ) + + output_mask = output_mask & node_ownership_mask( + graph.n_node, + graph.n_local, + atype.shape[0], + ) + if self.pair_excl is not None: + keep = self.pair_excl.build_edge_exclude_mask( + graph.edge_index, atype_clamped + ) + graph = dataclasses.replace( + graph, + edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype), + ) + if self.atom_excl is not None: + output_mask = xp.logical_and( + output_mask, + self.atom_excl.build_type_exclude_mask(atype_clamped), + ) + return graph, atype_clamped, output_mask + def forward_common_atomic_graph( self, graph: "NeighborGraph", @@ -352,18 +388,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)) - if self.pair_excl is not None: - keep = self.pair_excl.build_edge_exclude_mask( - graph.edge_index, atype_clamped - ) - graph = dataclasses.replace( - graph, - edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype), - ) + graph, atype_clamped, output_mask = self._prepare_graph_inputs(graph, atype) ret_dict = self.forward_atomic_graph( graph, atype_clamped, @@ -371,7 +396,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 _finalize_atomic_ret( self, ret_dict: dict, atom_mask: Array, atype: Array @@ -399,10 +424,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 3008daac84..09ae8c6948 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -431,10 +431,11 @@ 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. Remaining ineligible configs - (``exclude_types``, and compressed descriptors) fall back to the legacy - dense path, so those models keep working unchanged. + and transformer attention (``attn_layer >= 0``) with concat or strip + type-embedding. Geo-compressed strip models + (``geo_compress=True``, ``attn_layer == 0``) are also graph-eligible; + tebd-only compression (``geo_compress=False``) and ``exclude_types`` + stay on the legacy dense path. Eligibility does NOT imply numerical interchangeability with the dense route for every config: with ``smooth_type_embedding=True`` @@ -442,10 +443,13 @@ def uses_graph_lower(self) -> bool: differs from the dense lower by up to ~1e-4 (see the Notes of :meth:`call_graph`). """ - # compressed descriptors have no graph kernel (geo/tebd tabulation is - # dense-only); keep them on the legacy dense path. if self.compress: - return False + 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 + ) # exclude_types stays dense (graph exclusion is owned elsewhere); strip is # now graph-eligible (per-edge factorized embedding, no neighbor coupling). return ( @@ -1756,7 +1760,7 @@ def call_graph( ----- Known limitations: - ``tebd_input_mode`` in {"concat", "strip"}; compressed descriptors stay dense; - - ``exclude_types`` is not yet supported and raises (lands in a later PR). + - ``exclude_types`` is not supported and raises. """ from deepmd.dpmodel.utils.neighbor_graph import ( edge_env_mat, @@ -1768,10 +1772,7 @@ def call_graph( f"graph path does not support tebd_input_mode={self.tebd_input_mode!r}" ) if self.exclude_types: - raise NotImplementedError( - "graph path does not yet apply exclude_types (NeighborGraph PR-A); " - "type exclusion lands in a later PR" - ) + raise NotImplementedError("graph path does not support exclude_types") if type_embedding is None: raise ValueError("type_embedding is required for the graph path") xp = array_api_compat.array_namespace(graph.edge_vec) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 3933ec2533..6b1823d7e4 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -641,9 +641,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 @@ -651,18 +651,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 ------- @@ -676,6 +673,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 @@ -714,9 +712,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 @@ -724,18 +722,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 ca524991d2..83c3469b72 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. """ @@ -24,6 +25,10 @@ build_neighbor_graph, from_dense_quartet, ) +from .csr import ( + build_edge_csr, + canonicalize_neighbor_graph, +) from .derivatives import ( edge_force_virial, ) @@ -37,6 +42,7 @@ GraphLayout, NeighborGraph, frame_id_from_n_node, + node_ownership_mask, node_validity_mask, pad_and_guard_angles, pad_and_guard_edges, @@ -59,8 +65,10 @@ "angle_to_node_sum", "attach_angles", "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", @@ -68,6 +76,7 @@ "from_dense_quartet", "graph_angle_cos", "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 3b00ee6fac..42b81b7982 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py @@ -42,6 +42,9 @@ def build_neighbor_graph_ase( box: Array | None, rcut: float, layout: GraphLayout | None = None, + *, + with_csr: bool = False, + canonicalize: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using ASE's O(N) cell-list search. @@ -66,6 +69,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``. Returns ------- @@ -137,5 +146,14 @@ def _to_cpu_numpy(x: Any) -> np.ndarray: S_all, nframe_all = S_all[keep], nframe_all[keep] return 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, + with_csr=with_csr, + canonicalize=canonicalize, ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index 71ca699e1b..8f449cd4fc 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -39,6 +39,9 @@ import array_api_compat +from .csr import ( + build_edge_csr, +) from .graph import ( GraphLayout, NeighborGraph, @@ -57,6 +60,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). @@ -98,6 +104,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 ------- @@ -108,6 +120,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 @@ -148,11 +161,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_row_ptr, + source_order, + ) = 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_row_ptr=source_row_ptr, + source_order=source_order, + destination_sorted=canonicalize, ) else: # COMPACT: drop invalid slots via nonzero (dynamic shape -> eager only, @@ -189,11 +229,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_row_ptr, + source_order, + ) = 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_row_ptr=source_row_ptr, + source_order=source_order, + destination_sorted=canonicalize, ) @@ -203,6 +270,9 @@ def build_neighbor_graph( box: Array | None, rcut: float, layout: GraphLayout | None = None, + *, + with_csr: bool = False, + canonicalize: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph DIRECTLY from coordinates (``dense`` search). @@ -236,6 +306,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``. """ from deepmd.dpmodel.utils.nlist import ( extend_coord_with_ghosts, @@ -246,6 +322,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] @@ -298,9 +375,36 @@ def build_neighbor_graph( 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_row_ptr, + source_order, + ) = 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_row_ptr=source_row_ptr, + source_order=source_order, + destination_sorted=canonicalize, ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/csr.py b/deepmd/dpmodel/utils/neighbor_graph/csr.py new file mode 100644 index 0000000000..21900a0d61 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/csr.py @@ -0,0 +1,190 @@ +# 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_row_ptr : Array + Source 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``. + + 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_row_ptr, + source_order, + ) + + +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. + """ + if 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_row_ptr, + source_order, + ) = build_edge_csr( + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + n_nodes, + canonicalize=True, + ) + 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_row_ptr=source_row_ptr, + source_order=source_order, + destination_sorted=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..c3e5d37e64 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_row_ptr, + source_order, + ) = 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_row_ptr=source_row_ptr, + source_order=source_order, + destination_sorted=canonicalize, ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index 67f7d53845..4ca176bf50 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, @@ -33,7 +34,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 @@ -50,6 +54,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_row_ptr: Array | None = field(default=None, kw_only=True) + """(N + 1,) int64 CSR offsets into ``source_order``.""" + source_order: Array | None = field(default=None, kw_only=True) + """(E,) source-grouped edge permutation; same dtype as edge_index.""" + destination_sorted: bool = field(default=False, kw_only=True) + """Whether the payload is destination-major and destination_order is identity.""" @dataclass @@ -218,6 +232,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 node_validity_mask(n_node: Array, n_total: int) -> Array: """Derive the (n_total,) real-vs-padding node mask from per-frame counts. 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..913c7b4ef9 100644 --- a/deepmd/entrypoints/convert_backend.py +++ b/deepmd/entrypoints/convert_backend.py @@ -36,19 +36,23 @@ 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. + # Pass the optional flags a backend's deserialize hook accepts, probed by + # signature so scripts stay backend-agnostic. ``lower_kind="auto"`` lets the + # pt_expt hook pick the deploy-optimal lowering (the graph lower with its + # fused operator for a graph-eligible .pt2, the dense nlist lower otherwise); + # ``do_atomic_virial`` forwards ``--atomic-virial``. 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/graph_compress.py b/deepmd/kernels/cuda/dpa1/graph_compress.py new file mode 100644 index 0000000000..348f3b8fad --- /dev/null +++ b/deepmd/kernels/cuda/dpa1/graph_compress.py @@ -0,0 +1,995 @@ +# 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) + +_registered = False + + +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 _registered, _cpu_library + if _registered 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") + _registered = True + + +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_row_ptr is None + or graph.source_order 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_row_ptr, + graph.source_order, + 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..eef68ecd8e --- /dev/null +++ b/deepmd/kernels/cuda/dpa1/graph_descriptor.py @@ -0,0 +1,763 @@ +# 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 + +_registered = False + + +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 _registered, _cpu_library + if _registered 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") + _registered = True + + +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..d164a3bb88 --- /dev/null +++ b/deepmd/kernels/cuda/dpa1/graph_energy_force.py @@ -0,0 +1,433 @@ +# 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", +] + +_registered = False + + +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_row_ptr: torch.Tensor, + source_order: 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_row_ptr: torch.Tensor, + source_order: 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_row_ptr, + source_order, + 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 _registered, _cpu_library + if _registered 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") + _registered = True + + +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_row_ptr is None + or graph.source_order 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_row_ptr.contiguous(), + graph.source_order.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..13c35e92a9 --- /dev/null +++ b/deepmd/kernels/cuda/edge_force_virial.py @@ -0,0 +1,178 @@ +# 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__ = [ + "edge_force_virial", + "ensure_registered", + "op_available", +] + +_registered = False + + +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 _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_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]: + 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_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]: + 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 + + +_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 _registered, _cpu_library + if _registered or not op_available(): + return + torch.library.register_fake("deepmd::edge_force_virial")(_fake) + _cpu_library = torch.library.Library("deepmd", "IMPL") + _cpu_library.impl("edge_force_virial", _cpu, "CPU") + _registered = True + + +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_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 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_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..bcc04f4262 --- /dev/null +++ b/deepmd/kernels/cuda/graph_fitting.py @@ -0,0 +1,312 @@ +# 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", +] + +_registered = False + + +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 _registered, _cpu_library + if _registered 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") + _registered = True + + +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..99c4630519 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_row_ptr", + "source_order", +) + + +# ====================================================================== +# 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..420a71d9ae 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.train.training 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="graph", + ) + 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/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 f03b6facfb..192df59a20 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1748,11 +1748,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_row_ptr, source_order, + 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, @@ -1767,8 +1770,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( @@ -1780,8 +1783,35 @@ 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) + destination_order_t = torch.as_tensor( + graph.destination_order, + device=DEVICE, + ) + destination_row_ptr_t = torch.as_tensor( + graph.destination_row_ptr, + dtype=torch.int64, + device=DEVICE, + ) + source_row_ptr_t = torch.as_tensor( + graph.source_row_ptr, + dtype=torch.int64, + device=DEVICE, + ) + source_order_t = torch.as_tensor( + graph.source_order, + device=DEVICE, + ) fparam_t, aparam_t = self._prepare_optional_lower_inputs( fparam, aparam, nframes, natoms, DEVICE @@ -1791,9 +1821,14 @@ def _eval_model_graph( 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_row_ptr_t, + source_order_t, fparam_t, aparam_t, charge_spin_t, @@ -1828,7 +1863,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 if method == "dense": @@ -1836,16 +1873,26 @@ def _build_eval_graph( build_neighbor_graph, ) - return build_neighbor_graph(coord_input, atom_types, box_input, self._rcut) - if method == "ase": + graph = build_neighbor_graph( + coord_input, + atom_types, + box_input, + self._rcut, + canonicalize=True, + ) + elif method == "ase": from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph_ase, ) - return build_neighbor_graph_ase( - coord_input, atom_types, box_input, self._rcut + graph = build_neighbor_graph_ase( + coord_input, + atom_types, + box_input, + self._rcut, + canonicalize=True, ) - if method in ("vesin", "nv"): + elif method in ("vesin", "nv"): cc = torch.as_tensor(coord_input, dtype=torch.float64, device=device) aa = torch.as_tensor( np.asarray(atom_types), dtype=torch.int64, device=device @@ -1860,16 +1907,24 @@ def _build_eval_graph( build_neighbor_graph_vesin, ) - return build_neighbor_graph_vesin(cc, aa, bb, self._rcut) - from deepmd.pt_expt.utils.nv_graph_builder import ( - build_neighbor_graph_nv, + graph = build_neighbor_graph_vesin( + cc, aa, bb, self._rcut, canonicalize=True + ) + else: + from deepmd.pt_expt.utils.nv_graph_builder import ( + build_neighbor_graph_nv, + ) + + graph = build_neighbor_graph_nv( + cc, aa, bb, self._rcut, canonicalize=True + ) + else: + raise ValueError( + f"unknown neighbor_graph_method {method!r}; " + "use 'dense', 'ase', 'vesin', or 'nv'" ) - return build_neighbor_graph_nv(cc, aa, bb, self._rcut) - raise ValueError( - f"unknown neighbor_graph_method {method!r}; " - "use 'dense', 'ase', 'vesin', or 'nv'" - ) + return graph def _get_output_shape( self, odef: OutputVariableDef, nframes: int, natoms: int diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 653f323404..4f9ff77c08 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_row_ptr: torch.Tensor | None = None, + source_order: 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,6 +106,37 @@ def edge_energy_deriv( create_graph=create_graph, retain_graph=True, ) + 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_row_ptr is not None + and source_order 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_row_ptr, + source_order, + n_node, + n_cap, + do_atomic_virial, + ) + return force, (atom_virial if do_atomic_virial else None), virial force, atom_virial, virial = edge_force_virial( g_e, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity ) @@ -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_row_ptr, + graph.source_order, 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..cd33b3c12b 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -280,13 +280,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_row_ptr: torch.Tensor, + source_order: 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 +304,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. + destination_row_ptr, source_row_ptr + (N + 1,) destination/source CSR offsets. + source_order + (E,) source-grouped edge permutation. + 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 +334,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_row_ptr, + source_order, 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_row_ptr, + source_order, fparam=fparam, aparam=aparam, charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, + destination_sorted=destination_sorted, **make_fx_kwargs, ) @@ -348,9 +371,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_row_ptr: torch.Tensor, + source_order: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, @@ -358,9 +386,14 @@ def fn( model_ret = traced( atype, n_node, + n_local, edge_index, edge_vec, edge_mask, + destination_order, + destination_row_ptr, + source_row_ptr, + source_order, fparam, aparam, charge_spin, @@ -374,5 +407,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_row_ptr, + source_order, + fparam, + aparam, + charge_spin, ) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index ae2e83eada..c4a20b98b7 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_row_ptr: torch.Tensor | None = None, + source_order: 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_row_ptr + (N + 1,) source CSR offsets into ``source_order``. + source_order + (E,) edge permutation grouped by source node. + 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_row_ptr=source_row_ptr, + source_order=source_order, + 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,22 +567,51 @@ def _call_common_graph( "graph lower (e.g. dpa1 attn_layer=0)" ) rcut = self.get_rcut() + with_csr = ( + not self.training + and cuda_infer_level() >= 1 + and bool(getattr(descriptor, "geo_compress", False)) + ) if method == "dense": - ng = build_neighbor_graph(cc, atype, bb, rcut) + ng = build_neighbor_graph( + cc, + atype, + bb, + rcut, + with_csr=with_csr, + ) elif method == "ase": - ng = build_neighbor_graph_ase(cc, atype, bb, rcut) + ng = build_neighbor_graph_ase( + cc, + atype, + bb, + rcut, + with_csr=with_csr, + ) 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) + ng = build_neighbor_graph_vesin( + cc, + atype, + bb, + rcut, + with_csr=with_csr, + ) 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) + ng = build_neighbor_graph_nv( + cc, + atype, + bb, + rcut, + with_csr=with_csr, + ) else: raise ValueError( f"unknown neighbor_graph_method {method!r}; " @@ -494,9 +622,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_row_ptr, + ng.source_order, + destination_sorted=ng.destination_sorted, do_atomic_virial=do_atomic_virial, fparam=fp, aparam=ap, @@ -663,13 +797,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_row_ptr: torch.Tensor, + source_order: 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`` @@ -678,15 +818,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. + destination_row_ptr, source_row_ptr + (N + 1,) destination/source CSR offsets. + source_order + (E,) source-grouped edge permutation. + 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 @@ -697,18 +849,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_row_ptr, + source_order, 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_row_ptr, + source_order, + 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_row_ptr: torch.Tensor, + source_order: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, @@ -719,9 +888,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_row_ptr, + source_order, + destination_sorted=destination_sorted, do_atomic_virial=do_atomic_virial, fparam=fparam, aparam=aparam, @@ -731,9 +906,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_row_ptr, + source_order, fparam, aparam, charge_spin, @@ -766,13 +946,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 497e15713d..4794ee8dac 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -739,9 +739,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_row_ptr, + s_source_order, s_fparam, s_aparam, s_charge_spin, @@ -750,9 +755,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_row_ptr: torch.Tensor, + source_order: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, @@ -778,9 +788,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_row_ptr, + source_order, do_atomic_virial=False, fparam=fparam, aparam=aparam, @@ -813,9 +828,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_row_ptr, + s_source_order, s_fparam, s_aparam, s_charge_spin, @@ -1215,9 +1235,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_row_ptr, + ng.source_order, fparam, aparam, charge_spin, 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..f001328e3c --- /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_row_ptr: torch.Tensor, + source_order: 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 06a30c7dd4..5e9013f3cd 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -204,6 +204,9 @@ def build_neighbor_graph_nv( box: Any | None, rcut: float, layout: GraphLayout | None = None, + *, + with_csr: bool = False, + canonicalize: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using nvalchemiops' GPU cell list. @@ -219,6 +222,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``. Returns ------- @@ -237,6 +246,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 @@ -248,7 +260,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). @@ -256,8 +277,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 @@ -276,5 +301,14 @@ def build_neighbor_graph_nv( shift, frame_idx = shift[keep], frame_idx[keep] return 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, + with_csr=with_csr, + canonicalize=canonicalize, ) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 9317e87f4d..f9725ceddc 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_row_ptr, source_order, 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,13 +477,21 @@ 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_row_ptr, + graph.source_order, fparam, aparam, charge_spin, @@ -476,23 +503,25 @@ def _build_graph_dynamic_shapes( ) -> 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_row_ptr, source_order, + 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 +529,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: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) + {0: nedge_dim}, # source_order: (E,) {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,6 +652,42 @@ 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 == "graph" + 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" ) -> dict: @@ -738,6 +808,7 @@ def _probe_has_message_passing(obj: object) -> bool | None: # "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["graph_edge_dtype"] = _graph_edge_dtype(model, lower_kind) return meta @@ -804,6 +875,61 @@ 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.model import ( + BaseModel, + ) + from deepmd.pt_expt.train.training import ( + _model_uses_graph_lower, + ) + + model = BaseModel.deserialize(data["model"]) + return ( + "graph" + if _model_uses_graph_lower(model) and _supports_graph_export(model) + else "nlist" + ) + + def deserialize_to_file( model_file: str, data: dict, @@ -832,25 +958,43 @@ 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 == "graph": + 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( @@ -923,14 +1067,25 @@ 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) - # 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. + # Graph-form export uses the ghost-free NeighborGraph schema with a dynamic + # edge axis and an energy-model output contract. if lower_kind == "graph": import math @@ -942,7 +1097,7 @@ 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"): raise NotImplementedError( @@ -951,7 +1106,7 @@ def _trace_and_export( "requires an energy model" ) - # The edge axis is DYNAMIC (B2.0): the AOTI artifact accepts any edge + # 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 @@ -960,23 +1115,35 @@ def _trace_and_export( 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. + # make_fx traces on CPU. Conditioning inputs retain the model-agnostic + # float64 ABI; compressed DPA1 graph geometry enters directly in its + # float32 compute precision, as recorded in metadata. + 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"), ) ( atype_g, n_node_g, + n_local_g, edge_index_g, edge_vec_g, edge_mask_g, + destination_order_g, + destination_row_ptr_g, + source_row_ptr_g, + source_order_g, fparam_g, aparam_g, charge_spin_g, @@ -986,22 +1153,33 @@ def _trace_and_export( traced = model.forward_lower_graph_exportable( atype_g, n_node_g, + n_local_g, edge_index_g, edge_vec_g, edge_mask_g, + destination_order_g, + destination_row_ptr_g, + source_row_ptr_g, + source_order_g, fparam=fparam_g, aparam=aparam_g, do_atomic_virial=do_atomic_virial, charge_spin=charge_spin_g, + destination_sorted=True, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) sample_out = traced( atype_g, n_node_g, + n_local_g, edge_index_g, edge_vec_g, edge_mask_g, + destination_order_g, + destination_row_ptr_g, + source_row_ptr_g, + source_order_g, fparam_g, aparam_g, charge_spin_g, @@ -1020,11 +1198,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": @@ -1035,7 +1211,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. @@ -1143,6 +1319,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( @@ -1213,6 +1390,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, @@ -1353,21 +1531,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 @@ -1385,16 +1578,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 3f86fc7b75..ed8d3d5bca 100644 --- a/deepmd/pt_expt/utils/vesin_graph_builder.py +++ b/deepmd/pt_expt/utils/vesin_graph_builder.py @@ -89,11 +89,37 @@ def build_neighbor_graph_vesin( box: Any | None, rcut: float, layout: GraphLayout | None = None, + *, + with_csr: bool = False, + canonicalize: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using vesin.torch's O(N) cell list. Mirrors :func:`deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph_ase` but runs on the input tensor's device via ``vesin.torch``. + + Parameters + ---------- + 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``. + + Returns + ------- + NeighborGraph + The carry-all local-atom graph. """ if not is_vesin_torch_available(): raise ImportError( @@ -115,7 +141,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 = [], [], [], [] @@ -163,5 +198,14 @@ 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. return 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, + with_csr=with_csr, + canonicalize=canonicalize, ) diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 39a0dff883..751d3bee58 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -200,15 +200,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/pyproject.toml b/pyproject.toml index 4f3a707977..556d77434b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -305,7 +305,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_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 4a197cb064..6e19cb897a 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, @@ -324,7 +325,26 @@ class DeepPotBackend : public DeepBaseModelBackend { const int* d_edge_index, const double* d_edge_vec, const int nloc, - const int nedge); + 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 bool supports_device_edge_inference() const; + virtual bool uses_fp32_edge_vectors() const; }; /** @@ -683,7 +703,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 @@ -712,6 +733,68 @@ 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 Runtime frame parameters, or empty for the model default. + * @param[in] aparam 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; + int dim_chg_spin() const; protected: diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 6c30fae991..693633d0c4 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -291,24 +291,33 @@ class DeepPotPTExpt : public DeepPotBackend { const bool atomic) 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 Frame parameters, or empty for the stored default. + * @param[in] aparam 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 +327,54 @@ 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; + /** + * @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; 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); bool inited; int ntypes; int dfparam; @@ -339,6 +393,7 @@ 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 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 +401,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; std::unique_ptr with_comm_tempfile_; std::unique_ptr with_comm_loader; @@ -431,28 +482,38 @@ 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_row_ptr, source_order, + * [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_row_ptr Source CSR offsets ``(N + 1,)``. + * @param[in] source_order Source-grouped edge permutation ``(E,)``. */ - 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_row_ptr, + const torch::Tensor& source_order, + const torch::Tensor& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin); /** * @brief Run the with-comm .pt2 artifact with comm tensors appended. @@ -475,7 +536,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 02c25aa047..e02ca34175 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -369,28 +369,101 @@ 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_row_ptr; + torch::Tensor 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. @@ -403,7 +476,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). @@ -427,7 +500,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); @@ -441,7 +514,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 @@ -451,6 +524,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. @@ -459,6 +533,7 @@ inline GraphTensorPack buildGraphTensors( pack.atype = torch::from_blob(atype_nodes.data(), {n_node_count}, int_options) .clone() .to(device); + canonicalizeGraphPayload(pack, n_node_count); return pack; } @@ -488,7 +563,7 @@ inline GraphTensorPack buildGraphTensors( * 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 @@ -510,8 +585,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..4b4c45923d 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -572,7 +572,11 @@ void DeepPotBackend::compute_edges_gpu(double* d_atom_energy, const int* d_edge_index, const double* d_edge_vec, const int nloc, - const int nedge) { + 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; @@ -580,14 +584,53 @@ void DeepPotBackend::compute_edges_gpu(double* d_atom_energy, (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."); +} + +bool DeepPotBackend::uses_fp32_edge_vectors() const { return false; } + +bool DeepPotBackend::supports_device_edge_inference() const { return false; } + void DeepPot::compute_edges_gpu(double* d_atom_energy, double* d_force, double* d_atom_virial, @@ -597,14 +640,57 @@ 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. + // Parameter-free path: no runtime fparam/aparam, so the model defaults apply. + compute_edges_gpu(d_atom_energy, d_force, d_atom_virial, d_coord, d_atype, + d_edge_index, d_edge_vec, nloc, nedge, {}, {}, 0, nullptr); +} + +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); + d_edge_index, d_edge_vec, nloc, nedge, fparam, aparam, + nall_nodes, comm_nlist); +} + +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(); } int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); } diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 9088762fcb..16b39d0d2e 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -2,6 +2,8 @@ #include "DeepPotPTExpt.h" #if defined(BUILD_PYTORCH) && BUILD_PT_EXPT +#include +#include #include #include @@ -11,6 +13,7 @@ #include #include #include +#include #include "SimulationRegion.h" #include "common.h" @@ -26,6 +29,73 @@ using deepmd::ptexpt::read_zip_entry; using namespace deepmd; +namespace { + +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 +137,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); @@ -161,6 +231,16 @@ void DeepPotPTExpt::init(const std::string& model, lower_input_is_edge_ = false; lower_input_is_graph_ = 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"; + } type_map.clear(); for (const auto& v : metadata["type_map"].as_array()) { @@ -179,26 +259,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(); if (has_comm_artifact_) { @@ -295,16 +367,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_row_ptr, + const torch::Tensor& source_order, 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_row_ptr, source_order}; if (dfparam > 0) { inputs.push_back(fparam); } @@ -328,12 +406,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( @@ -372,12 +447,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( @@ -488,10 +560,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], @@ -503,32 +574,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( @@ -559,10 +613,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++) { @@ -609,13 +662,9 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, 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. + // 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, @@ -657,18 +706,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. @@ -719,14 +757,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 @@ -808,7 +842,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 @@ -840,9 +874,8 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, } else if (lower_input_is_graph_) { 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.) @@ -859,27 +892,42 @@ 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}); - flat_outputs = - run_model_graph(node_atype, n_node_tensor, edge_tensors.edge_index, - edge_tensors.edge_vec, edge_tensors.edge_mask, - fparam_tensor, aparam_tensor, charge_spin_tensor); + 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 = edge_tensors.edge_mask; + canonicalizeGraphPayload(graph_pack, n_node_count); + 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_row_ptr, graph_pack.source_order, fparam_tensor, + graph_aparam, charge_spin_tensor); } else { flat_outputs = run_model(coord_Tensor, atype_Tensor, firstneigh_tensor, mapping_tensor, fparam_tensor, aparam_tensor, @@ -1040,7 +1088,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, @@ -1184,18 +1241,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. @@ -1239,9 +1286,14 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, edge_tensors.edge_index_ext, edge_tensors.edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor); } else if (lower_input_is_graph_) { + const at::Tensor graph_edge_vec = + graph_edge_fp32_ ? graph_tensors.edge_vec.to(torch::kFloat32) + : graph_tensors.edge_vec; flat_outputs = run_model_graph( - graph_tensors.atype, graph_tensors.n_node, graph_tensors.edge_index, - graph_tensors.edge_vec, graph_tensors.edge_mask, fparam_tensor, + graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, + graph_tensors.edge_index, graph_edge_vec, graph_tensors.edge_mask, + graph_tensors.destination_order, graph_tensors.destination_row_ptr, + graph_tensors.source_row_ptr, graph_tensors.source_order, fparam_tensor, aparam_tensor, charge_spin_tensor); } else { flat_outputs = @@ -1325,8 +1377,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. @@ -1360,16 +1419,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()) { @@ -1522,9 +1583,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. @@ -1560,16 +1635,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()) { @@ -1748,28 +1825,62 @@ 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) { + 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 " @@ -1782,14 +1893,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 = @@ -1798,51 +1915,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()) { @@ -1858,6 +1994,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; @@ -1866,35 +2035,168 @@ 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 = edge_mask; + 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_row_ptr, graph_pack.source_order) = + 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_row_ptr, graph_pack.source_order, 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); + } + at::cuda::getCurrentCUDAStream().synchronize(); }); } + +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); +} + +bool DeepPotPTExpt::uses_fp32_edge_vectors() const { + return lower_input_is_graph_ && graph_edge_fp32_; +} + +bool DeepPotPTExpt::supports_device_edge_inference() const { + return lower_input_is_edge_ || lower_input_is_graph_; +} #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..03413844c8 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,10 @@ 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()); double ener; std::vector force, virial; dp.compute(ener, force, virial, coord, atype, box); @@ -130,10 +129,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 +190,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 +238,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 +256,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 +340,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..ff88cbcea7 100644 --- a/source/api_cc/tests/test_deeppot_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_ptexpt.cc @@ -136,6 +136,7 @@ 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()); double ener; std::vector force, virial; dp.compute(ener, force, virial, coord, atype, box); @@ -644,7 +645,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..546e43a6ba 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()); @@ -163,6 +162,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 +212,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/builtin.cmake b/source/lmp/builtin.cmake index e051e5c24a..42862ea918 100644 --- a/source/lmp/builtin.cmake +++ b/source/lmp/builtin.cmake @@ -68,7 +68,23 @@ target_sources( ${LAMMPS_SOURCE_DIR}/KSPACE/remap_wrap.cpp ${LAMMPS_SOURCE_DIR}/EXTRA-FIX/fix_ttm.cpp # for ttm ) -target_link_libraries(lammps PUBLIC DeePMD::deepmd_c) +if(PKG_KOKKOS) + # The Kokkos device pair style (pair_style deepmd/kk) calls + # DeepPot::compute_edges_gpu, which is exposed only by the C++ API; link the + # C++ library and select the C++ API headers in pair_deepmd.h / pair_base.h. + target_link_libraries(lammps PUBLIC DeePMD::deepmd_cc) + target_compile_definitions(lammps PRIVATE DP_USE_CXX_API) + # deepmd_cc exports an interface precompiled header (common.h/device.h) that + # lives under /include/deepmd; expose that directory (PUBLIC, so the + # dependent lmp executable inherits it) to resolve the PCH. + get_target_property(_deepmd_inc DeePMD::deepmd_cc + INTERFACE_INCLUDE_DIRECTORIES) + foreach(_inc IN LISTS _deepmd_inc) + target_include_directories(lammps PUBLIC "${_inc}/deepmd") + endforeach() +else() + target_link_libraries(lammps PUBLIC DeePMD::deepmd_c) +endif() target_include_directories( lammps PRIVATE ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_LIST_DIR} ${LAMMPS_SOURCE_DIR}/KSPACE ${LAMMPS_SOURCE_DIR}/EXTRA-FIX) 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..7895fa4b88 --- /dev/null +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -0,0 +1,767 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#ifdef LMP_KOKKOS +#include "pair_deepmd_kokkos.h" + +#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), + 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(); + +#ifndef DP_USE_CXX_API + error->all(FLERR, + "pair style deepmd/kk requires the C++ API build of DeePMD " + "(compute_edges_gpu is unavailable through the C API)."); +#endif + // 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(); + // 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); +} + +#ifdef DP_USE_CXX_API +template +int PairDeepMDKokkos::build_edges_device() { + const int nlocal = atom->nlocal; + const int nall = atom->nlocal + atom->nghost; + + // === Ghost fold / model-node compaction (host, amortized to rebuilds) === + // The graph indexes model nodes. ``loc2model`` maps each atom to its node + // (-1 for a NULL/virtual type); ``model2loc`` inverts it with local atoms + // ordered first so the energy sum over [0, nloc_model) covers exactly the + // local atoms. Single rank folds a ghost neighbour onto the local atom it + // images (``owner``, so nnode_model == nloc_model); multi-rank keeps real + // ghosts as their own extended nodes (nnode_model > nloc_model) and the + // LAMMPS reverse communication later folds their forces back to the owners. + // atom->map / tag / type are host structures, rebuilt only at ago == 0. + 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(); + } + + // === 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 | TYPE_MASK); + auto x = atomKK->k_x.template view(); + auto type = atomKK->k_type.template view(); + + const double cut = cutoff; + const double cutsq = cut * cut; + const bool multi = multi_rank; + auto owner = d_owner; + auto type_map = d_type_map; + 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); + } + if ((int)d_model_type.extent(0) < nnode_model) { + d_model_type = Kokkos::View("deepmd/kk:model_type", nall); + } + auto edge_offset = d_edge_offset; + auto model_type = d_model_type; + + // Model type per node (local first, then ghost for the extended set). + 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); + }); + + // === 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; +} +#endif + +template +void PairDeepMDKokkos::compute(int eflag, int vflag) { +#ifdef DP_USE_CXX_API + 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(); + } + const int 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 (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 (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, 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, 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; + } + } +#else + PairDeepMD::compute(eflag, vflag); +#endif +} + +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..11336da721 --- /dev/null +++ b/source/lmp/pair_deepmd_kokkos.h @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// The Kokkos device pair style is compiled only when the LAMMPS KOKKOS package +// is enabled; without it the translation unit is empty, so the GLOB that adds +// this file to the build is harmless for non-Kokkos LAMMPS. +#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 "kokkos_base.h" +#include "kokkos_type.h" +#include "neigh_list_kokkos.h" +#include "pair_deepmd.h" + +namespace LAMMPS_NS { + +// 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. + int build_edges_device(); + + 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 + // 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 + + // 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 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/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 3bb34e622d..4f07c19cc5 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 and 13.x CCCL fail to compile CUB/Thrust with -arch=all. + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.9" AND CUDAToolkit_VERSION + VERSION_LESS "14.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..134025fd40 --- /dev/null +++ b/source/op/pt/dpa1_graph_compress.cu @@ -0,0 +1,1512 @@ +// 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]); +} + +__device__ __forceinline__ bool edge_is_active(long edge, + const bool* edge_mask) { + return edge_mask[edge]; +} + +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.data_ptr(), + destination_order.data_ptr(), + 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), + 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.data_ptr(), + destination_order.data_ptr(), + 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), + 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.data_ptr(), + 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 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) { \ + 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, + 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) { \ + 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, 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, + 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()); +} + +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); +} 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..c006b81daf --- /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_row_ptr, + torch::Tensor source_order, + 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_row_ptr, + source_order, 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_row_ptr, Tensor source_order, 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..7b6106282a --- /dev/null +++ b/source/op/pt/edge_force_virial.cu @@ -0,0 +1,405 @@ +// 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 long* __restrict__ source_row_ptr, + const index_t* __restrict__ source_order, + 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 = static_cast(destination_order[position]); + if (!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]) { + 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_row_ptr, + const torch::Tensor& source_order, + 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.data_ptr(), + destination_order.data_ptr(), + destination_row_ptr.data_ptr(), source_row_ptr.data_ptr(), + source_order.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"); +} + +} // 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_row_ptr, source_order}; +} + +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_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(); + const long frame_count = n_node_per_frame.size(0); + TORCH_CHECK(edge_gradient.is_cuda() && edge_vec.is_cuda() && + edge_mask.is_cuda() && destination_order.is_cuda() && + destination_row_ptr.is_cuda() && source_row_ptr.is_cuda() && + source_order.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_row_ptr.is_contiguous() && + source_order.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"); + + 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_row_ptr, + source_order, 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_row_ptr, + source_order, frame_row_ptr, force, node_virial, virial_partial, + virial, stream); + } + }); + return {force, atom_virial, 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_row_ptr, Tensor source_order)"); + 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_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("edge_force_virial", torch::kCUDA, &edge_force_virial); +} + +TORCH_LIBRARY_IMPL(deepmd, Autograd, library) { + library.impl("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..a3cff383b5 --- /dev/null +++ b/source/op/pt/graph_fitting.cu @@ -0,0 +1,358 @@ +// 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 + +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); +} + +// dpre = 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)); + } + auto scratch = torch::empty({n_node, width_max}, f32); + // 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(), scratch.data_ptr(), + (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, scratch.data_ptr(), + 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. +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; + const long input_width = ws[0].size(0); + auto d_x = torch::empty({n_node, input_width}, saved.options()); + // Guard the empty system before the division by ``n_node`` below. + if (n_node == 0) { + return d_x; + } + 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 dpre = torch::empty({n_node, width_max}, f32); + 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; + backward_epilogue_kernel<<>>( + n_node * dout / 4, dh.data_ptr(), adot, dpre.data_ptr()); + FITTING_CHECK_LAUNCH("graph_fitting backward layer"); + 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; + } + gemm_nt(stream, dpre.data_ptr(), ws[l].data_ptr(), out, + (int)n_node, din, dout, beta); + if (l > 0) { + std::swap(dh, dh_next); + } + } + 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..43c812e996 --- /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_row_ptr, + torch::Tensor source_order, + torch::Tensor n_node_per_frame, + c10::SymInt node_capacity, + bool want_atom_virial); 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_block.py b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py index fb01caec55..bb8c59bd0a 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py @@ -90,9 +90,6 @@ def test_block_graph_equals_dense_any_sel(self, sel, type_one_side) -> None: atol=1e-12, ) - # attn_layer > 0 is supported since NeighborGraph PR-D; parity is covered - # by test_dpa1_graph_attention_parity.py (the fail-fast test was removed). - def test_exclude_types_raises(self) -> None: """The graph block kernel fail-fasts for exclude_types (not yet applied).""" # the graph path does not yet apply type exclusion; it must fail-fast 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 6632d82894..27561cf318 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py @@ -125,8 +125,8 @@ def test_ineligible_config_falls_back_to_dense(self, kwargs) -> 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..dd4f695b53 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_row_ptr, + source_order, + ) = 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_row_ptr, + source_order, + ) = 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_row_ptr=np.array([0, 1, 2], dtype=np.int64), + source_order=np.array([0, 1], 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 9ba25c0ccb..66a7647e0d 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -100,6 +100,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_row_ptr) + self.assertIsNone(generic.source_order) + self.assertFalse(with_csr.destination_sorted) + self.assertIsNotNone(with_csr.destination_order) + self.assertIsNotNone(with_csr.destination_row_ptr) + self.assertIsNotNone(with_csr.source_row_ptr) + self.assertIsNotNone(with_csr.source_order) + 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) 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..e2d1a49bae --- /dev/null +++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py @@ -0,0 +1,1721 @@ +# 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_row_ptr=( + graph.source_row_ptr.cpu() if graph.source_row_ptr is not None else None + ), + source_order=( + graph.source_order.cpu() if graph.source_order 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_row_ptr=graph.source_row_ptr.cpu(), + source_order=graph.source_order.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_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_row_ptr, + source_order, + ) = 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_row_ptr=source_row_ptr, + source_order=source_order, + 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_row_ptr=None, + source_order=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_row_ptr=torch.zeros(1, dtype=torch.int64, device=self.device), + source_order=order, + 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_row_ptr=torch.zeros(1, dtype=torch.int64), + source_order=order, + 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_row_ptr, + graph.source_order, + 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_row_ptr=None, + source_order=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_row_ptr, + graph.source_order, + ) + + 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_row_ptr, + graph.source_order, + 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_row_ptr, + graph.source_order, + 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_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_row_ptr, + source_order, + ) = 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_row_ptr=source_row_ptr, + source_order=source_order, + 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_row_ptr, source_order = ( + 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_row_ptr, + src_order, + ) = 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_row_ptr), + move(src_order), + move(n_node), + total, + ) + + def _reference( + self, + g_e, + edge_vec, + edge_index, + mask, + dst_order, + dst_row_ptr, + src_row_ptr, + src_order, + 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_row_ptr, + src_order, + 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_row_ptr, + src_order, + 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_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, + row_ptr, + order, + 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_row_ptr, + source_order, + ) = 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_row_ptr.cuda(), + source_order.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 11c49cdca2..5b870f0b56 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -103,8 +103,8 @@ def _make_model(self, attn_layer: int = 0, smooth: bool = False) -> EnergyModel: 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", @@ -169,14 +169,14 @@ 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 def test_force_virial_parity_vs_legacy(self, periodic, do_av, attn_layer) -> None: """Graph lower energy/force/virial/atom_virial == legacy dense lower on the SAME neighbor set (regime-1 graph from from_dense_quartet). attn_layer=2 exercises graph attention through model-level autograd - (smooth=False: exact carry-all parity regime, NeighborGraph PR-D). + (smooth=False: exact carry-all parity regime). """ model = self._make_model(attn_layer=attn_layer) model.eval() @@ -205,6 +205,7 @@ def test_force_virial_parity_vs_legacy(self, periodic, do_av, attn_layer) -> Non graph = model.forward_common_lower_graph( atype_local, ng.n_node, + ng.n_node, ng.edge_index, ng.edge_vec, ng.edge_mask, @@ -269,23 +270,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_row_ptr, + source_order, + fp, + ap, + cs, + ) = sample traced = model.forward_lower_graph_exportable( atype, n_node, + n_local, ei, ev, em, + destination_order, + destination_row_ptr, + source_row_ptr, + source_order, 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_row_ptr, + source_order, + 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_row_ptr, + source_order, + 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) @@ -302,7 +349,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. @@ -383,6 +430,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..4b7e426ca2 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_row_ptr, + g.source_order, + ) 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_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..57648489b4 --- /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_row_ptr, + source_order, + ) = 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_row_ptr, + source_order, + ) + + +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_row_ptr, + source_order, + ) = _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_row_ptr, + source_order, + 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_row_ptr, + source_order, + ) = _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_row_ptr, + source_order, + 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) From a89ac04eaeb3cac1d0109fda5d591deaf8fc9722 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 11 Jul 2026 18:41:59 +0800 Subject: [PATCH 2/9] perf(pt/dpa1): add compact canonical graph deployment Introduce a source-only canonical graph ABI for eligible geometrically compressed DPA1 models, allowing LAMMPS Kokkos to pass cutoff-compacted dual-CSR topology directly into AOTInductor without rebuilding graph metadata each step. - replace the generic ten-tensor deployment graph with a dedicated eight-tensor ABI containing atom types, owned/total node counts, int64 source indices, FP32 edge vectors, destination/source row pointers, and source order - derive destinations implicitly from canonical CSR rows and remove deployment edge masks, destination indices, destination order, index casts, guard concatenation, and bridge-side CSR construction - add compact compressed-descriptor forward/backward and force/virial operators while sharing the existing CUDA kernels with the generic graph path - build destination and source CSR views directly in the Kokkos cutoff producer with capacity-backed persistent storage and checked size_t/int64 accounting - preserve owned/halo energy semantics, NULL-type compaction, empty-rank handling, centroid virial reverse communication, and standard C++ inference - select the compact ABI automatically when pt_expt compress exports an eligible DPA1 model, while retaining the full NeighborGraph ABI for attention, exclusions, masked topology, and unsupported configurations - reduce fitting workspaces by writing GEMMs directly into activation slots and converting backward derivatives in place - add dynamic export, E=0/1 guards, CUDA parity, C++ packing, DeepEval, single/multi-rank, and large-system regression coverage At 129,168 atoms the compact path reduces whole-step time by 16.8%, 14.8%, and 10.3% for S, M, and L. The M process GPU-memory peak drops by 1,034 MiB, and the S/M/L OOM ceilings increase to 9.52M, 7.53M, and 5.51M atoms. --- .../dpmodel/atomic_model/base_atomic_model.py | 47 +- deepmd/kernels/cuda/dpa1/canonical.py | 456 ++++++++++++++++++ deepmd/kernels/cuda/edge_force_virial.py | 142 ++++++ deepmd/pt_expt/entrypoints/compress.py | 2 +- deepmd/pt_expt/infer/deep_eval.py | 78 ++- deepmd/pt_expt/model/ener_model.py | 162 +++++++ deepmd/pt_expt/utils/canonical_graph.py | 203 ++++++++ deepmd/pt_expt/utils/serialization.py | 268 ++++++---- source/api_cc/include/DeepPot.h | 32 ++ source/api_cc/include/DeepPotPTExpt.h | 37 ++ source/api_cc/include/commonPT.h | 43 ++ source/api_cc/src/DeepPot.cc | 55 +++ source/api_cc/src/DeepPotPTExpt.cc | 189 +++++++- .../tests/test_deeppot_dpa1_graph_ptexpt.cc | 1 + source/api_cc/tests/test_deeppot_ptexpt.cc | 1 + .../api_cc/tests/test_neighbor_list_data.cc | 25 + source/lmp/pair_deepmd_kokkos.cpp | 284 +++++++++-- source/lmp/pair_deepmd_kokkos.h | 20 + source/op/pt/dpa1_graph_compress.cu | 120 ++++- source/op/pt/edge_force_virial.cu | 167 +++++-- source/op/pt/graph_fitting.cu | 58 ++- .../pt_expt/descriptor/test_dpa1_cuda.py | 295 +++++++++++ .../pt_expt/utils/test_canonical_graph.py | 49 ++ 23 files changed, 2470 insertions(+), 264 deletions(-) create mode 100644 deepmd/kernels/cuda/dpa1/canonical.py create mode 100644 deepmd/pt_expt/utils/canonical_graph.py create mode 100644 source/tests/pt_expt/utils/test_canonical_graph.py diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index f72ff27380..28fd0400b1 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -313,27 +313,49 @@ 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_inputs( + def _prepare_graph_nodes( self, - graph: "NeighborGraph", + n_node: Array, + n_local: Array | None, atype: Array, - ) -> tuple["NeighborGraph", Array, Array]: - """Apply graph masks shared by standard and fused atomic forwards.""" - 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 + 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 graph.n_local is not None: + if n_local is not None: from deepmd.dpmodel.utils.neighbor_graph import ( node_ownership_mask, ) output_mask = output_mask & node_ownership_mask( - graph.n_node, - graph.n_local, + 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.""" + xp = array_api_compat.array_namespace(graph.edge_vec) + atype_clamped, output_mask = self._prepare_graph_nodes( + graph.n_node, + graph.n_local, + atype, + graph.edge_vec, + ) if self.pair_excl is not None: keep = self.pair_excl.build_edge_exclude_mask( graph.edge_index, atype_clamped @@ -342,11 +364,6 @@ def _prepare_graph_inputs( graph, edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype), ) - if self.atom_excl is not None: - output_mask = xp.logical_and( - output_mask, - self.atom_excl.build_type_exclude_mask(atype_clamped), - ) return graph, atype_clamped, output_mask def forward_common_atomic_graph( diff --git a/deepmd/kernels/cuda/dpa1/canonical.py b/deepmd/kernels/cuda/dpa1/canonical.py new file mode 100644 index 0000000000..5fc52d2482 --- /dev/null +++ b/deepmd/kernels/cuda/dpa1/canonical.py @@ -0,0 +1,456 @@ +# 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", +] + +_registered = False +_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 _registered, _cpu_library + if _registered 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", + ) + _registered = True + + +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/edge_force_virial.py b/deepmd/kernels/cuda/edge_force_virial.py index 13c35e92a9..5b54302e29 100644 --- a/deepmd/kernels/cuda/edge_force_virial.py +++ b/deepmd/kernels/cuda/edge_force_virial.py @@ -30,6 +30,8 @@ import torch __all__ = [ + "canonical_edge_force_virial", + "canonical_op_available", "edge_force_virial", "ensure_registered", "op_available", @@ -44,6 +46,12 @@ def op_available() -> bool: 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, @@ -65,6 +73,25 @@ def _fake( ) +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, @@ -93,6 +120,64 @@ def _cpu( 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_row_ptr, + source_order, + n_node_per_frame, + node_capacity, + want_atom_virial, + ) + + _cpu_library: torch.library.Library | None = None @@ -105,8 +190,18 @@ def ensure_registered() -> None: if _registered 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", + ) _registered = True @@ -176,3 +271,50 @@ def edge_force_virial( 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/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 420a71d9ae..ba2949e097 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -127,7 +127,7 @@ def enable_compression( deserialize_to_file( output, {"model": compressed_model_dict, "model_def_script": model_def_script}, - lower_kind="graph", + lower_kind="auto", ) else: log.info("Re-exporting compressed model...") diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 192df59a20..8872049fb0 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -184,10 +184,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 " @@ -1508,7 +1507,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 ) @@ -1813,26 +1812,57 @@ def _eval_model_graph( device=DEVICE, ) - 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) + 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, + ) - 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_row_ptr_t, - source_order_t, - fparam_t, - aparam_t, - charge_spin_t, - ) + 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_row_ptr=source_row_ptr_t, + source_order=source_order_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_row_ptr_t, + source_order_t, + fparam_t, + aparam_t, + charge_spin_t, + ) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index cd33b3c12b..2710edb4c6 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, diff --git a/deepmd/pt_expt/utils/canonical_graph.py b/deepmd/pt_expt/utils/canonical_graph.py new file mode 100644 index 0000000000..ce449009d1 --- /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_row_ptr is None + or graph.source_order 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/serialization.py b/deepmd/pt_expt/utils/serialization.py index f9725ceddc..83832d1577 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -498,6 +498,90 @@ def build_synthetic_graph_inputs( ) +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_row_ptr, + source_order, + _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_row_ptr=source_row_ptr, + source_order=source_order, + 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: @@ -665,7 +749,7 @@ def _graph_edge_dtype(model: torch.nn.Module, lower_kind: str) -> str: descriptor_block = getattr(descriptor, "se_atten", None) statistics = getattr(descriptor_block, "mean", None) if ( - lower_kind == "graph" + lower_kind in ("graph", "dpa1_canonical") and bool(getattr(descriptor, "geo_compress", False)) and isinstance(statistics, torch.Tensor) and statistics.dtype == torch.float32 @@ -689,7 +773,9 @@ def _supports_graph_export(model: torch.nn.Module) -> bool: 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. @@ -807,7 +893,7 @@ 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) return meta @@ -923,11 +1009,13 @@ def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str: ) model = BaseModel.deserialize(data["model"]) - return ( - "graph" - if _model_uses_graph_lower(model) and _supports_graph_export(model) - else "nlist" - ) + 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( @@ -981,7 +1069,7 @@ def deserialize_to_file( # DP_CUDA_INFER >= 2 so the analytic backward and CSR scatter remain custom # operators, while the per-atom virial is mandatory for the LAMMPS Kokkos # consumer. - if lower_kind == "graph": + if lower_kind in ("graph", "dpa1_canonical"): do_atomic_virial = True ctx: contextlib.AbstractContextManager = _cuda_infer_at_least_2() else: @@ -989,11 +1077,19 @@ def deserialize_to_file( with ctx: if model_file.endswith(".pt2"): _deserialize_to_file_pt2( - model_file, data, model_json_override, do_atomic_virial, lower_kind + 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 + model_file, + data, + model_json_override, + do_atomic_virial, + lower_kind, ) @@ -1082,11 +1178,14 @@ def _trace_and_export( 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, + ) - # Graph-form export uses the ghost-free NeighborGraph schema with a dynamic - # edge axis and an energy-model output contract. - 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) @@ -1099,94 +1198,71 @@ def _trace_and_export( "graph-form .pt2 export does not support the with-comm artifact " "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 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. Conditioning inputs retain the model-agnostic - # float64 ABI; compressed DPA1 graph geometry enters directly in its - # float32 compute precision, as recorded in metadata. - 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"), - ) - - ( - atype_g, - n_node_g, - n_local_g, - edge_index_g, - edge_vec_g, - edge_mask_g, - destination_order_g, - destination_row_ptr_g, - source_row_ptr_g, - source_order_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, - n_local_g, - edge_index_g, - edge_vec_g, - edge_mask_g, - destination_order_g, - destination_row_ptr_g, - source_row_ptr_g, - source_order_g, - fparam=fparam_g, - aparam=aparam_g, - do_atomic_virial=do_atomic_virial, - charge_spin=charge_spin_g, - destination_sorted=True, - tracing_mode="symbolic", - _allow_non_fake_inputs=True, - ) - sample_out = traced( - atype_g, - n_node_g, - n_local_g, - edge_index_g, - edge_vec_g, - edge_mask_g, - destination_order_g, - destination_row_ptr_g, - source_row_ptr_g, - source_order_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, diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 6e19cb897a..1cab1bc5a3 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -343,8 +343,22 @@ class DeepPotBackend : public DeepBaseModelBackend { 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; }; /** @@ -795,6 +809,24 @@ class DeepPot : public DeepBaseModel { */ 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 693633d0c4..c89758beb9 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -332,6 +332,18 @@ class DeepPotPTExpt : public DeepPotBackend { 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. * @@ -359,6 +371,7 @@ class DeepPotPTExpt : public DeepPotBackend { * @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 @@ -375,6 +388,19 @@ class DeepPotPTExpt : public DeepPotBackend { 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; @@ -393,6 +419,7 @@ 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) @@ -515,6 +542,16 @@ class DeepPotPTExpt : public DeepPotBackend { 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. * diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index e02ca34175..0cfdf9d982 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -379,6 +379,49 @@ struct GraphTensorPack { torch::Tensor source_order; }; +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. * diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 4b4c45923d..a7d55c9644 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -627,10 +627,42 @@ void DeepPotBackend::compute_edges_gpu(double* d_atom_energy, "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, @@ -685,6 +717,25 @@ void DeepPot::compute_edges_gpu(double* d_atom_energy, 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(); } @@ -693,6 +744,10 @@ 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 16b39d0d2e..45c6f9bc8d 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -227,9 +227,11 @@ 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")) { @@ -241,6 +243,12 @@ void DeepPotPTExpt::init(const std::string& model, } 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(); for (const auto& v : metadata["type_map"].as_array()) { @@ -395,6 +403,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, @@ -661,7 +682,7 @@ 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_) { + } 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. @@ -871,7 +892,7 @@ 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, which violates the exported @@ -922,12 +943,20 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, : edge_tensors.edge_vec; graph_pack.edge_mask = edge_tensors.edge_mask; canonicalizeGraphPayload(graph_pack, n_node_count); - 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_row_ptr, graph_pack.source_order, fparam_tensor, - graph_aparam, charge_spin_tensor); + 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_row_ptr, graph_pack.source_order, fparam_tensor, + graph_aparam, charge_spin_tensor); + } } else { flat_outputs = run_model(coord_Tensor, atype_Tensor, firstneigh_tensor, mapping_tensor, fparam_tensor, aparam_tensor, @@ -939,7 +968,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. @@ -1203,7 +1232,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 = @@ -1285,16 +1314,26 @@ 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_) { const at::Tensor graph_edge_vec = graph_edge_fp32_ ? graph_tensors.edge_vec.to(torch::kFloat32) : graph_tensors.edge_vec; - flat_outputs = run_model_graph( - graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, - graph_tensors.edge_index, graph_edge_vec, graph_tensors.edge_mask, - graph_tensors.destination_order, graph_tensors.destination_row_ptr, - graph_tensors.source_row_ptr, graph_tensors.source_order, fparam_tensor, - aparam_tensor, charge_spin_tensor); + graph_tensors.edge_vec = graph_edge_vec; + 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_row_ptr, + graph_tensors.source_order, fparam_tensor, aparam_tensor, + charge_spin_tensor); + } } else { flat_outputs = run_model(coord_Tensor, atype_Tensor, nlist_tensor, mapping_tensor, @@ -1305,7 +1344,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 @@ -1839,6 +1878,11 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, 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( @@ -2156,6 +2200,85 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, }); } +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); + at::cuda::getCurrentCUDAStream().synchronize(); + }); +} + void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, double* d_force, double* d_atom_virial, @@ -2192,11 +2315,37 @@ void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy, 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_ && graph_edge_fp32_; + 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_; + 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/tests/test_deeppot_dpa1_graph_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa1_graph_ptexpt.cc index 03413844c8..1a22e80038 100644 --- a/source/api_cc/tests/test_deeppot_dpa1_graph_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_dpa1_graph_ptexpt.cc @@ -113,6 +113,7 @@ TYPED_TEST(TestInferDpa1GraphPtExpt, cpu_build_nlist) { 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); diff --git a/source/api_cc/tests/test_deeppot_ptexpt.cc b/source/api_cc/tests/test_deeppot_ptexpt.cc index ff88cbcea7..b4f12d1453 100644 --- a/source/api_cc/tests/test_deeppot_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_ptexpt.cc @@ -137,6 +137,7 @@ TYPED_TEST(TestInferDeepPotAPtExpt, cpu_build_nlist) { 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); diff --git a/source/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc index 546e43a6ba..487643e421 100644 --- a/source/api_cc/tests/test_neighbor_list_data.cc +++ b/source/api_cc/tests/test_neighbor_list_data.cc @@ -138,6 +138,31 @@ TEST(TestNeighborListData, RoundTripWithEmptyRows) { EXPECT_EQ(out.numneigh[3], 1); } +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_row_ptr = torch::tensor({0, 1}, torch::kInt64); + graph.source_order = torch::tensor({0, 1, 2}, 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); +} + #ifdef BUILD_PYTORCH TEST(TestEdgeTensorPack, CreateEdgeTensorsUsesRowCenters) { const torch::Device device(torch::kCPU); diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index 7895fa4b88..e813e48fc3 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -2,9 +2,11 @@ #ifdef LMP_KOKKOS #include "pair_deepmd_kokkos.h" +#include #include #include #include +#include #include "atom.h" #include "atom_kokkos.h" @@ -34,6 +36,7 @@ PairDeepMDKokkos::PairDeepMDKokkos(LAMMPS* lmp) nnode_model(0), edge_capacity(0), edge_vec_fp32(false), + canonical_graph(false), device_path_ok(false), reverse_virial(false), reverse_used_host(false) { @@ -220,6 +223,7 @@ void PairDeepMDKokkos::init_style() { 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. @@ -243,19 +247,10 @@ void PairDeepMDKokkos::init_style() { #ifdef DP_USE_CXX_API template -int PairDeepMDKokkos::build_edges_device() { +void PairDeepMDKokkos::prepare_model_nodes() { const int nlocal = atom->nlocal; const int nall = atom->nlocal + atom->nghost; - // === Ghost fold / model-node compaction (host, amortized to rebuilds) === - // The graph indexes model nodes. ``loc2model`` maps each atom to its node - // (-1 for a NULL/virtual type); ``model2loc`` inverts it with local atoms - // ordered first so the energy sum over [0, nloc_model) covers exactly the - // local atoms. Single rank folds a ghost neighbour onto the local atom it - // images (``owner``, so nnode_model == nloc_model); multi-rank keeps real - // ghosts as their own extended nodes (nnode_model > nloc_model) and the - // LAMMPS reverse communication later folds their forces back to the owners. - // atom->map / tag / type are host structures, rebuilt only at ago == 0. 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); @@ -304,6 +299,40 @@ int PairDeepMDKokkos::build_edges_device() { 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); @@ -312,15 +341,13 @@ int PairDeepMDKokkos::build_edges_device() { auto d_neighbors = k_list->d_neighbors; auto d_ilist = k_list->d_ilist; - atomKK->sync(execution_space, X_MASK | TYPE_MASK); + atomKK->sync(execution_space, X_MASK); auto x = atomKK->k_x.template view(); - auto type = atomKK->k_type.template view(); const double cut = cutoff; const double cutsq = cut * cut; const bool multi = multi_rank; auto owner = d_owner; - auto type_map = d_type_map; auto loc2model = d_loc2model; auto model2loc = d_model2loc; @@ -328,18 +355,7 @@ int PairDeepMDKokkos::build_edges_device() { d_edge_offset = Kokkos::View( "deepmd/kk:edge_offset", nlocal + 1); } - if ((int)d_model_type.extent(0) < nnode_model) { - d_model_type = Kokkos::View("deepmd/kk:model_type", nall); - } auto edge_offset = d_edge_offset; - auto model_type = d_model_type; - - // Model type per node (local first, then ghost for the extended set). - 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); - }); // === 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 @@ -482,6 +498,186 @@ int PairDeepMDKokkos::build_edges_device() { } return nedge; } + +template +std::int64_t PairDeepMDKokkos::build_canonical_edges_device( + CompactCanonicalGraphWorkspace& workspace) { + prepare_model_nodes(); + + auto* k_list = static_cast*>(list); + const int inum = k_list->inum; + auto d_numneigh = k_list->d_numneigh; + auto d_neighbors = k_list->d_neighbors; + auto d_ilist = k_list->d_ilist; + atomKK->sync(execution_space, X_MASK); + auto x = atomKK->k_x.template view(); + auto owner = d_owner; + auto loc2model = d_loc2model; + const bool multi = multi_rank; + const double cutsq = cutoff * cutoff; + const double inv_dist = 1.0 / dist_unit_cvt_factor; + const int node_count_int = nnode_model; + const std::size_t node_count = static_cast(node_count_int); + + if (workspace.destination_row_ptr.extent(0) < node_count + 1) { + workspace.destination_row_ptr = Kokkos::View( + "deepmd/kk:canonical_destination_row_ptr", node_count + 1); + workspace.source_counts = Kokkos::View( + "deepmd/kk:canonical_source_counts", node_count); + workspace.source_row_ptr = Kokkos::View( + "deepmd/kk:canonical_source_row_ptr", node_count + 1); + workspace.source_cursor = Kokkos::View( + "deepmd/kk:canonical_source_cursor", node_count); + } + Kokkos::deep_copy(workspace.destination_row_ptr, std::int64_t{0}); + Kokkos::deep_copy(workspace.source_counts, std::int64_t{0}); + if (node_count_int == 0) { + return 0; + } + auto destination_row_ptr = workspace.destination_row_ptr; + auto source_counts = workspace.source_counts; + + Kokkos::parallel_for( + "deepmd/kk:canonical_count", Kokkos::RangePolicy(0, inum), + KOKKOS_LAMBDA(const int ii) { + const int i = d_ilist(ii); + const int mi = loc2model(i); + if (mi < 0) { + return; + } + const double xi = x(i, 0); + const double yi = x(i, 1); + const double zi = x(i, 2); + const int jnum = d_numneigh(i); + std::int64_t count = 0; + for (int jj = 0; jj < jnum; ++jj) { + int j = d_neighbors(i, jj) & NEIGHMASK; + const int mj = loc2model(multi ? j : owner(j)); + if (mj < 0) { + continue; + } + const double dx = x(j, 0) - xi; + const double dy = x(j, 1) - yi; + const double dz = x(j, 2) - zi; + if (dx * dx + dy * dy + dz * dz < cutsq) { + ++count; + } + } + destination_row_ptr(mi) = count; + }); + + Kokkos::parallel_scan( + "deepmd/kk:canonical_destination_scan", + Kokkos::RangePolicy(0, node_count_int), + KOKKOS_LAMBDA(const int node, std::int64_t& update, const bool final) { + const std::int64_t count = destination_row_ptr(node); + if (final) { + destination_row_ptr(node) = update; + } + update += count; + if (final && node == node_count_int - 1) { + destination_row_ptr(node_count_int) = update; + } + }); + std::int64_t edge_count = 0; + Kokkos::deep_copy(edge_count, Kokkos::subview(workspace.destination_row_ptr, + node_count_int)); + const std::int64_t storage_count = std::max(edge_count, 2); + const std::size_t required = static_cast(storage_count); + if (workspace.edge_capacity < required) { + const std::size_t slack = required / 8 + 64; + if (required > std::numeric_limits::max() - slack) { + error->one(FLERR, "Compact DPA1 graph capacity overflows size_t"); + } + workspace.edge_capacity = required + slack; + workspace.source = Kokkos::View( + "deepmd/kk:canonical_source", workspace.edge_capacity); + workspace.edge_vec = Kokkos::View( + "deepmd/kk:canonical_edge_vec", workspace.edge_capacity * 3); + workspace.source_order = Kokkos::View( + "deepmd/kk:canonical_source_order", workspace.edge_capacity); + } + + auto source = workspace.source; + auto edge_vec = workspace.edge_vec; + Kokkos::parallel_for( + "deepmd/kk:canonical_fill", Kokkos::RangePolicy(0, inum), + KOKKOS_LAMBDA(const int ii) { + const int i = d_ilist(ii); + const int mi = loc2model(i); + if (mi < 0) { + return; + } + const double xi = x(i, 0); + const double yi = x(i, 1); + const double zi = x(i, 2); + const int jnum = d_numneigh(i); + std::int64_t edge = destination_row_ptr(mi); + for (int jj = 0; jj < jnum; ++jj) { + int j = d_neighbors(i, jj) & NEIGHMASK; + const int mj = loc2model(multi ? j : owner(j)); + if (mj < 0) { + continue; + } + const double dx = x(j, 0) - xi; + const double dy = x(j, 1) - yi; + const double dz = x(j, 2) - zi; + if (dx * dx + dy * dy + dz * dz < cutsq) { + source(edge) = static_cast(mj); + edge_vec(3 * edge + 0) = static_cast(dx * inv_dist); + edge_vec(3 * edge + 1) = static_cast(dy * inv_dist); + edge_vec(3 * edge + 2) = static_cast(dz * inv_dist); + Kokkos::atomic_fetch_add(&source_counts(mj), std::int64_t{1}); + ++edge; + } + } + }); + + auto source_row_ptr = workspace.source_row_ptr; + Kokkos::parallel_scan( + "deepmd/kk:canonical_source_scan", + Kokkos::RangePolicy(0, node_count_int), + KOKKOS_LAMBDA(const int node, std::int64_t& update, const bool final) { + const std::int64_t count = source_counts(node); + if (final) { + source_row_ptr(node) = update; + } + update += count; + if (final && node == node_count_int - 1) { + source_row_ptr(node_count_int) = update; + } + }); + Kokkos::deep_copy( + workspace.source_cursor, + Kokkos::subview(workspace.source_row_ptr, + std::make_pair(std::int64_t{0}, static_cast( + node_count_int)))); + auto source_cursor = workspace.source_cursor; + auto source_order = workspace.source_order; + Kokkos::parallel_for( + "deepmd/kk:canonical_source_scatter", + Kokkos::RangePolicy>( + 0, edge_count), + KOKKOS_LAMBDA(const std::int64_t edge) { + const auto position = + Kokkos::atomic_fetch_add(&source_cursor(source(edge)), 1LL); + source_order(position) = edge; + }); + if (storage_count > edge_count) { + Kokkos::parallel_for( + "deepmd/kk:canonical_guards", + Kokkos::RangePolicy>( + edge_count, storage_count), + KOKKOS_LAMBDA(const std::int64_t edge) { + source(edge) = 0; + edge_vec(3 * edge + 0) = 0.0f; + edge_vec(3 * edge + 1) = 0.0f; + edge_vec(3 * edge + 2) = 0.0f; + source_order(edge) = edge; + }); + } + return edge_count; +} #endif template @@ -508,7 +704,12 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { memoryKK->create_kokkos(k_eatom, eatom, maxeatom, "deepmd/kk:eatom"); d_eatom = k_eatom.template view(); } - const int nedge = build_edges_device(); + 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) @@ -592,7 +793,7 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { comm_ptr = &comm_list; } - if (nloc_m > 0 || comm_ptr != nullptr) { + 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 @@ -605,18 +806,27 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // them; edge-input models consume them). const double* coord_ptr = has_null_types ? d_coord_model.data() : x.data(); try { - 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, nedge, - fparam, aparam_step, nnode_m, comm_ptr); + 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, nedge, fparam, - aparam_step, nnode_m, comm_ptr); + deep_pot.compute_edges_gpu( + d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), + coord_ptr, d_model_type.data(), d_edge_index.data(), + d_edge_vec.data(), nloc_m, static_cast(nedge), fparam, + aparam_step, nnode_m, comm_ptr); } } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index 11336da721..1fdcf780c9 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -19,6 +19,7 @@ PairStyle(deepmd/kk/host,PairDeepMDKokkos); #ifndef LMP_PAIR_DEEPMD_KOKKOS_H #define LMP_PAIR_DEEPMD_KOKKOS_H +#include #include #include "kokkos_base.h" @@ -28,6 +29,18 @@ PairStyle(deepmd/kk/host,PairDeepMDKokkos); namespace LAMMPS_NS { +template +struct CompactCanonicalGraphWorkspace { + Kokkos::View source; + Kokkos::View edge_vec; + Kokkos::View destination_row_ptr; + Kokkos::View source_counts; + Kokkos::View source_row_ptr; + Kokkos::View source_cursor; + Kokkos::View source_order; + std::size_t edge_capacity = 0; +}; + // GPU-resident inference for exported ``.pt2`` models whose forward consumes // an explicit edge graph: both the graph-input form (a compact, unpadded // neighbor graph) and the edge-input form. Both are dispatched through @@ -74,13 +87,18 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // 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; @@ -108,6 +126,7 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { Kokkos::View d_edge_vec; // (3 * nedge) Kokkos::View d_edge_vec_float; // (3 * nedge), compressed graph ABI + CompactCanonicalGraphWorkspace canonical_workspace; // Model outputs on the device. Energy is per local atom; force and virial // span the model node set (up to ``nall`` under domain decomposition). @@ -124,6 +143,7 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { 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 diff --git a/source/op/pt/dpa1_graph_compress.cu b/source/op/pt/dpa1_graph_compress.cu index 134025fd40..cb9e141f33 100644 --- a/source/op/pt/dpa1_graph_compress.cu +++ b/source/op/pt/dpa1_graph_compress.cu @@ -333,7 +333,7 @@ __device__ __forceinline__ long edge_at_csr_position( __device__ __forceinline__ bool edge_is_active(long edge, const bool* edge_mask) { - return edge_mask[edge]; + return edge_mask == nullptr || edge_mask[edge]; } template @@ -940,8 +940,10 @@ void launch_forward_variant(long node_count, 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.data_ptr(), - destination_order.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(), @@ -1067,8 +1069,10 @@ void launch_backward_variant(long node_count, 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.data_ptr(), - destination_order.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(), @@ -1147,7 +1151,9 @@ void launch_backward(long node_count, launch(config, node_count); COMPRESS_CHECK_LAUNCH("dpa1_graph_compress backward"); zero_padding_kernel<<<1, kThreads, 0, stream>>>( - node_count, edge_count, destination_order.data_ptr(), + 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"); } @@ -1483,6 +1489,86 @@ torch::Tensor dpa1_graph_compress_backward( 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, " @@ -1509,4 +1595,26 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "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/edge_force_virial.cu b/source/op/pt/edge_force_virial.cu index 7b6106282a..1fabd9d371 100644 --- a/source/op/pt/edge_force_virial.cu +++ b/source/op/pt/edge_force_virial.cu @@ -79,8 +79,10 @@ __global__ void edge_force_virial_kernel( const long destination_end = destination_row_ptr[node + 1]; for (long position = destination_begin + lane; position < destination_end; position += 32) { - const long edge = static_cast(destination_order[position]); - if (!edge_mask[edge]) { + 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]; @@ -93,7 +95,7 @@ __global__ void edge_force_virial_kernel( for (long position = source_begin + lane; position < source_end; position += 32) { const long edge = static_cast(source_order[position]); - if (!edge_mask[edge]) { + if (edge_mask && !edge_mask[edge]) { continue; } const scalar_t gx = edge_gradient[edge * 3 + 0]; @@ -229,8 +231,10 @@ void launch_force_virial(long node_count, edge_force_virial_kernel <<>>( node_count, edge_gradient.data_ptr(), - edge_vec.data_ptr(), edge_mask.data_ptr(), - destination_order.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_row_ptr.data_ptr(), source_order.data_ptr(), force.data_ptr(), node_virial.data_ptr()); @@ -250,6 +254,62 @@ void launch_force_virial(long node_count, 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_row_ptr, + const torch::Tensor& source_order, + 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_row_ptr, + source_order, 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_row_ptr, + source_order, frame_row_ptr, force, node_virial, virial_partial, + virial, stream); + } + }); + return {force, atom_virial, virial}; +} + } // namespace std::tuple @@ -311,7 +371,6 @@ std::tuple edge_force_virial( c10::SymInt node_capacity, bool want_atom_virial) { const long node_count = node_capacity.expect_int(); - const long frame_count = n_node_per_frame.size(0); TORCH_CHECK(edge_gradient.is_cuda() && edge_vec.is_cuda() && edge_mask.is_cuda() && destination_order.is_cuda() && destination_row_ptr.is_cuda() && source_row_ptr.is_cuda() && @@ -338,49 +397,53 @@ std::tuple edge_force_virial( destination_order.scalar_type() == source_order.scalar_type(), "edge_force_virial: destination_order and source_order must have the " "same int32 or int64 dtype"); + return assemble_force_virial(node_count, edge_gradient, edge_vec, edge_mask, + destination_order, destination_row_ptr, + source_row_ptr, source_order, n_node_per_frame, + want_atom_virial); +} - 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(); +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"); - 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_row_ptr, - source_order, 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_row_ptr, - source_order, frame_row_ptr, force, node_virial, virial_partial, - virial, stream); - } - }); - return {force, atom_virial, virial}; + 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_row_ptr, source_order, n_node_per_frame, + want_atom_virial); } TORCH_LIBRARY_FRAGMENT(deepmd, library) { @@ -398,8 +461,18 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "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 index a3cff383b5..60724679c8 100644 --- a/source/op/pt/graph_fitting.cu +++ b/source/op/pt/graph_fitting.cu @@ -26,6 +26,7 @@ #include #include +#include #include namespace { @@ -175,7 +176,7 @@ __global__ void seed_kernel(long total4, make_float4(de * wv.x, de * wv.y, de * wv.z, de * wv.w); } -// dpre = dh * adot. +// Convert dh to dpre in place: dh *= adot. __global__ void backward_epilogue_kernel(long total4, const float* __restrict__ dh, const float* __restrict__ adot, @@ -229,7 +230,6 @@ std::tuple graph_fitting( for (int l = 0; l < n_layer; ++l) { width_max = std::max(width_max, (long)ws[l].size(1)); } - auto scratch = torch::empty({n_node, width_max}, f32); // 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). @@ -242,14 +242,13 @@ std::tuple graph_fitting( 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(), scratch.data_ptr(), - (int)n_node, dout, din); + 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, scratch.data_ptr(), + 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); @@ -276,11 +275,12 @@ std::tuple graph_fitting( // 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. -torch::Tensor graph_fitting_backward(torch::Tensor d_e, - torch::Tensor saved, - std::vector ws, - std::vector resnets, - torch::Tensor w_head) { +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); @@ -290,10 +290,15 @@ torch::Tensor graph_fitting_backward(torch::Tensor d_e, "the fitting widths"); const long n_node = saved.numel() / total_width; const long input_width = ws[0].size(0); - auto d_x = torch::empty({n_node, input_width}, saved.options()); + 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 d_x; + return; } auto stream = at::cuda::getCurrentCUDAStream(); const int n_layer = (int)ws.size(); @@ -308,7 +313,6 @@ torch::Tensor graph_fitting_backward(torch::Tensor d_e, for (int l = 0; l < n_layer; ++l) { width_max = std::max(width_max, (long)ws[l].size(1)); } - auto dpre = torch::empty({n_node, width_max}, f32); auto dh = torch::empty({n_node, width_max}, f32); auto dh_next = torch::empty({n_node, width_max}, f32); @@ -323,10 +327,6 @@ torch::Tensor graph_fitting_backward(torch::Tensor d_e, 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; - backward_epilogue_kernel<<>>( - n_node * dout / 4, dh.data_ptr(), adot, dpre.data_ptr()); - FITTING_CHECK_LAUNCH("graph_fitting backward layer"); float* out = l > 0 ? dh_next.data_ptr() : d_x.data_ptr(); const bool residual = resnets[l] && dout == din; float beta = 0.f; @@ -336,12 +336,34 @@ torch::Tensor graph_fitting_backward(torch::Tensor d_e, cudaMemcpyDeviceToDevice, stream); beta = 1.f; } - gemm_nt(stream, dpre.data_ptr(), ws[l].data_ptr(), out, + 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; } diff --git a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py index e2d1a49bae..a346d9a44e 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py @@ -568,6 +568,125 @@ def run(graph) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 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) @@ -1364,6 +1483,99 @@ def test_parity_vs_separate_ops(self) -> None: 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 ( @@ -1596,6 +1808,89 @@ def _assert_device_parity(self, device) -> None: 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_row_ptr, + source_order, + ) = 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_row_ptr=source_row_ptr, + source_order=source_order, + 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_row_ptr, + source_order, + 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 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..88a6cffc37 --- /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_row_ptr=row_ptr.clone(), + source_order=torch.arange(storage_count, dtype=torch.int64), + 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) From 3f13664a49c22ad74a3d79df2821701bd7457488 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 11 Jul 2026 19:33:17 +0800 Subject: [PATCH 3/9] perf(dpa1): specialize compact descriptor mask handling Instantiate mask-free compressed descriptor forward and backward kernels for cutoff-compacted canonical graphs, while retaining authoritative mask checks for generic cached, excluded, and permutation graph inputs. Include graph validity mode in the adaptive tuning key so masked and compact workloads select launch configurations independently. This improves compact model-only inference by about 1.4% without changing outputs or the generic graph contract. --- source/op/pt/dpa1_graph_compress.cu | 97 ++++++++++++++++++++--------- 1 file changed, 69 insertions(+), 28 deletions(-) diff --git a/source/op/pt/dpa1_graph_compress.cu b/source/op/pt/dpa1_graph_compress.cu index cb9e141f33..c06009b2a4 100644 --- a/source/op/pt/dpa1_graph_compress.cu +++ b/source/op/pt/dpa1_graph_compress.cu @@ -331,12 +331,20 @@ __device__ __forceinline__ long edge_at_csr_position( return static_cast(destination_order[position]); } +template __device__ __forceinline__ bool edge_is_active(long edge, const bool* edge_mask) { - return edge_mask == nullptr || edge_mask[edge]; + if constexpr (Masked) { + return edge_mask[edge]; + } + return true; } -template +template __global__ __launch_bounds__(kThreads, MinimumBlocks) void compressed_forward_kernel( long node_count, @@ -397,7 +405,7 @@ __launch_bounds__(kThreads, MinimumBlocks) void compressed_forward_kernel( 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 (!edge_is_active(edge, edge_mask)) { continue; } EdgeEnvironment environment{}; @@ -448,7 +456,7 @@ __launch_bounds__(kThreads, MinimumBlocks) void compressed_forward_kernel( 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 (!edge_is_active(edge, edge_mask)) { continue; } EdgeEnvironment environment{}; @@ -578,7 +586,11 @@ __launch_bounds__(kThreads, MinimumBlocks) void compressed_forward_kernel( } } -template +template __global__ __launch_bounds__(kThreads, MinimumBlocks) void compressed_backward_kernel( long node_count, @@ -723,7 +735,7 @@ __launch_bounds__(kThreads, MinimumBlocks) void compressed_backward_kernel( 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 (!edge_is_active(edge, edge_mask)) { if (half_lane == 0) { edge_gradient[edge * 3 + 0] = 0.0f; edge_gradient[edge * 3 + 1] = 0.0f; @@ -801,7 +813,7 @@ __launch_bounds__(kThreads, MinimumBlocks) void compressed_backward_kernel( 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 (!edge_is_active(edge, edge_mask)) { if (lane == 0) { edge_gradient[edge * 3 + 0] = 0.0f; edge_gradient[edge * 3 + 1] = 0.0f; @@ -897,7 +909,11 @@ __global__ void zero_padding_kernel( } } -template +template void launch_forward_variant(long node_count, int threads, int ntypes, @@ -934,7 +950,7 @@ void launch_forward_variant(long node_count, const int warps_per_block = threads / kWarpSize; const int blocks = static_cast((node_count + warps_per_block - 1) / warps_per_block); - compressed_forward_kernel + compressed_forward_kernel <<>>( node_count, ntypes, one_side, smooth, axis, concatenate_type_embedding, write_rotation, type_embedding_dim, rcut, @@ -952,7 +968,7 @@ void launch_forward_variant(long node_count, moment.data_ptr()); } -template +template void launch_forward(long node_count, int ntypes, bool one_side, @@ -995,7 +1011,8 @@ void launch_forward(long node_count, Canonical ? 1 : 0, static_cast(sizeof(index_t)), (one_side ? 1 : 0) | (smooth ? 2 : 0) | - (concatenate_type_embedding ? 4 : 0) | (write_rotation ? 8 : 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), @@ -1003,7 +1020,7 @@ void launch_forward(long node_count, }; const auto launch = [&](const LaunchConfig& config, long count) { if (config.resource == ResourcePolicy::kOccupancy) { - launch_forward_variant( + 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, @@ -1011,7 +1028,7 @@ void launch_forward(long node_count, destination_row_ptr, atype, type_embedding, average, inverse_stddev, table, gate_table, descriptor, rotation, moment, stream); } else { - launch_forward_variant( + 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, @@ -1026,7 +1043,11 @@ void launch_forward(long node_count, COMPRESS_CHECK_LAUNCH("dpa1_graph_compress forward"); } -template +template void launch_backward_variant(long node_count, int threads, long edge_count, @@ -1062,7 +1083,7 @@ void launch_backward_variant(long node_count, const int warps_per_block = threads / kWarpSize; const int blocks = static_cast((node_count + warps_per_block - 1) / warps_per_block); - compressed_backward_kernel + compressed_backward_kernel <<>>( node_count, edge_count, ntypes, one_side, smooth, axis, descriptor_stride, rcut, rcut_smooth, protection, inverse_neighbors, @@ -1079,7 +1100,7 @@ void launch_backward_variant(long node_count, edge_gradient.data_ptr()); } -template +template void launch_backward(long node_count, long edge_count, int ntypes, @@ -1121,7 +1142,7 @@ void launch_backward(long node_count, Canonical ? 1 : 0, static_cast(sizeof(index_t)), (one_side ? 1 : 0) | (smooth ? 2 : 0) | - (rotation_gradient != nullptr ? 8 : 0), + (rotation_gradient != nullptr ? 8 : 0) | (Masked ? 16 : 0), descriptor_stride, type_count_class(ntypes), workload_size_class(node_count, properties.multiProcessorCount), @@ -1129,7 +1150,7 @@ void launch_backward(long node_count, }; const auto launch = [&](const LaunchConfig& config, long count) { if (config.resource == ResourcePolicy::kOccupancy) { - launch_backward_variant( + 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, @@ -1137,7 +1158,7 @@ void launch_backward(long node_count, destination_order, destination_row_ptr, atype, average, inverse_stddev, table, gate_table, edge_gradient, stream); } else { - launch_backward_variant( + 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, @@ -1166,6 +1187,7 @@ void dispatch_forward(int width, bool smooth, int axis, bool canonical, + bool masked, bool concatenate_type_embedding, bool write_rotation, int type_embedding_dim, @@ -1195,8 +1217,17 @@ void dispatch_forward(int width, cudaStream_t stream) { #define DISPATCH_WIDTH(value) \ if (width == value) { \ - if (canonical) { \ - launch_forward( \ + 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, \ @@ -1205,7 +1236,7 @@ void dispatch_forward(int width, average, inverse_stddev, table, gate_table, descriptor, rotation, \ moment, stream); \ } else { \ - launch_forward( \ + 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, \ @@ -1235,6 +1266,7 @@ void dispatch_backward(int width, bool smooth, int axis, bool canonical, + bool masked, int descriptor_stride, float rcut, float rcut_smooth, @@ -1262,8 +1294,16 @@ void dispatch_backward(int width, cudaStream_t stream) { #define DISPATCH_WIDTH(value) \ if (width == value) { \ - if (canonical) { \ - launch_backward( \ + 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, \ @@ -1271,7 +1311,7 @@ void dispatch_backward(int width, destination_order, destination_row_ptr, atype, average, \ inverse_stddev, table, gate_table, edge_gradient, stream); \ } else { \ - launch_backward( \ + 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, \ @@ -1394,8 +1434,9 @@ std::tuple dpa1_graph_compress( using index_t = decltype(index_tag); dispatch_forward( width, node_count, ntypes, type_one_side != 0, smooth != 0, - static_cast(axis), canonical, concatenate_type_embedding != 0, - write_rotation != 0, type_embedding_dim, static_cast(rcut), + 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), @@ -1470,7 +1511,7 @@ torch::Tensor dpa1_graph_compress_backward( using index_t = decltype(index_tag); dispatch_backward( width, node_count, edge_count, ntypes, type_one_side != 0, smooth != 0, - static_cast(axis), canonical, + 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), From b1859a5f53c87e8244c379939d6ec49f2fa362c7 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 12 Jul 2026 13:39:59 +0800 Subject: [PATCH 4/9] fix(pt_expt): restore portable automatic graph export Keep convert-backend command-free for graph-eligible artifacts and isolate CUDA stream headers from ROCm builds. --- deepmd/entrypoints/convert_backend.py | 12 +----------- deepmd/main.py | 6 ------ deepmd/pt_expt/utils/serialization.py | 9 ++++----- source/api_cc/src/DeepPotPTExpt.cc | 6 ++++-- source/tests/test_convert_backend.py | 21 ++++----------------- 5 files changed, 13 insertions(+), 41 deletions(-) diff --git a/deepmd/entrypoints/convert_backend.py b/deepmd/entrypoints/convert_backend.py index 7b6af989c3..43cb901449 100644 --- a/deepmd/entrypoints/convert_backend.py +++ b/deepmd/entrypoints/convert_backend.py @@ -16,7 +16,6 @@ def convert_backend( INPUT: str, OUTPUT: str, atomic_virial: bool = False, - lower_kind: str = "nlist", **kwargs: Any, ) -> None: """Convert a model file from one backend to another. @@ -31,10 +30,6 @@ def convert_backend( If True, export .pt2/.pte models with per-atom virial correction. This adds ~2.5x inference cost. Default False. Silently ignored (with a warning) for backends that don't support the flag. - lower_kind : str - Lower input schema for output backends that support it. ``"nlist"`` is - the compatibility default; ``"graph"`` and ``"auto"`` are explicit - graph-lowering choices. """ inp_backend: Backend = Backend.detect_backend_by_model(INPUT)() out_backend: Backend = Backend.detect_backend_by_model(OUTPUT)() @@ -46,12 +41,7 @@ def convert_backend( sig = inspect.signature(out_hook) hook_kwargs: dict[str, Any] = {} if "lower_kind" in sig.parameters: - hook_kwargs["lower_kind"] = lower_kind - elif lower_kind != "nlist": - raise ValueError( - f"lower_kind={lower_kind!r} is not supported by output backend " - f"{out_backend.name}" - ) + hook_kwargs["lower_kind"] = "auto" if "do_atomic_virial" in sig.parameters: hook_kwargs["do_atomic_virial"] = atomic_virial elif atomic_virial: diff --git a/deepmd/main.py b/deepmd/main.py index 2813562926..42a03bc3f0 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -1001,12 +1001,6 @@ def main_parser() -> argparse.ArgumentParser: ) parser_convert_backend.add_argument("INPUT", help="The input model file.") parser_convert_backend.add_argument("OUTPUT", help="The output model file.") - parser_convert_backend.add_argument( - "--lower-kind", - choices=("nlist", "graph", "auto"), - default="nlist", - help="Lower input schema for pt_expt .pt2/.pte outputs.", - ) parser_convert_backend.add_argument( "--atomic-virial", action="store_true", diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 53ebd1cb84..0614d5527c 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1077,11 +1077,10 @@ def deserialize_to_file( (``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"`` (available through ``convert-backend --lower-kind auto``) - 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. + ``"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``. """ diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 26392cef9d..6689428b99 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -3,7 +3,7 @@ #if defined(BUILD_PYTORCH) && BUILD_PT_EXPT #include -#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) +#if defined(GOOGLE_CUDA) #include #endif #include @@ -34,8 +34,10 @@ using namespace deepmd; namespace { void synchronize_current_accelerator_stream() { -#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) +#if defined(GOOGLE_CUDA) at::cuda::getCurrentCUDAStream().synchronize(); +#elif defined(TENSORFLOW_USE_ROCM) + DPErrcheck(gpuDeviceSynchronize()); #else throw deepmd::deepmd_exception( "GPU-resident inference requires a GPU-enabled DeePMD-kit build."); diff --git a/source/tests/test_convert_backend.py b/source/tests/test_convert_backend.py index aea643e2b0..063caec868 100644 --- a/source/tests/test_convert_backend.py +++ b/source/tests/test_convert_backend.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Any, -) - import pytest from deepmd.backend.backend import ( @@ -13,19 +9,10 @@ ) -@pytest.mark.parametrize( - ("expected_lower_kind", "kwargs"), - [ - ("nlist", {}), - ("graph", {"lower_kind": "graph"}), - ], -) -def test_convert_backend_lower_kind( +def test_convert_backend_automatically_selects_lower_kind( monkeypatch: pytest.MonkeyPatch, - expected_lower_kind: str, - kwargs: dict[str, str], ) -> None: - captured: dict[str, Any] = {} + captured: dict[str, object] = {} class InputBackend: name = "input" @@ -57,7 +44,7 @@ def detect_backend(path: str) -> type[InputBackend] | type[OutputBackend]: monkeypatch.setattr(Backend, "detect_backend_by_model", detect_backend) - convert_backend(INPUT="model.input", OUTPUT="model.output", **kwargs) + convert_backend(INPUT="model.input", OUTPUT="model.output") - assert captured["lower_kind"] == expected_lower_kind + assert captured["lower_kind"] == "auto" assert captured["do_atomic_virial"] is False From b132c9c8b794afe50a5685ad81901103ea0f1751 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 12 Jul 2026 13:50:38 +0800 Subject: [PATCH 5/9] fix(ptexpt): avoid optional CUDA library headers Use the lightweight c10 stream API so CUDA builds do not require cuSPARSE headers while preserving current-stream synchronization. --- source/api_cc/src/DeepPotPTExpt.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 6689428b99..fdca92a306 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -4,7 +4,7 @@ #if defined(BUILD_PYTORCH) && BUILD_PT_EXPT #include #if defined(GOOGLE_CUDA) -#include +#include #endif #include #include @@ -35,7 +35,7 @@ namespace { void synchronize_current_accelerator_stream() { #if defined(GOOGLE_CUDA) - at::cuda::getCurrentCUDAStream().synchronize(); + c10::cuda::getCurrentCUDAStream().synchronize(); #elif defined(TENSORFLOW_USE_ROCM) DPErrcheck(gpuDeviceSynchronize()); #else From 410b5458e19757ac2e10e29571252c752462a088 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 12 Jul 2026 14:02:58 +0800 Subject: [PATCH 6/9] fix(ptexpt): synchronize through accelerator runtime Avoid optional PyTorch CUDA headers in portable builds while preserving the device-output completion contract. --- source/api_cc/src/DeepPotPTExpt.cc | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index fdca92a306..45a18da7fb 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -3,9 +3,6 @@ #if defined(BUILD_PYTORCH) && BUILD_PT_EXPT #include -#if defined(GOOGLE_CUDA) -#include -#endif #include #include @@ -34,9 +31,7 @@ using namespace deepmd; namespace { void synchronize_current_accelerator_stream() { -#if defined(GOOGLE_CUDA) - c10::cuda::getCurrentCUDAStream().synchronize(); -#elif defined(TENSORFLOW_USE_ROCM) +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) DPErrcheck(gpuDeviceSynchronize()); #else throw deepmd::deepmd_exception( From 254015486ca0696818d4519226633b207738c3eb Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 13 Jul 2026 12:29:43 +0800 Subject: [PATCH 7/9] refactor(graph): align dual-CSR field ordering Use parallel order/row-pointer pairs consistently across Python, export, CUDA, and C++ interfaces, and make the standard force fallback branch explicit. --- .../dpmodel/utils/neighbor_graph/builder.py | 8 +-- deepmd/dpmodel/utils/neighbor_graph/csr.py | 10 ++-- .../dpmodel/utils/neighbor_graph/from_ijs.py | 4 +- deepmd/dpmodel/utils/neighbor_graph/graph.py | 6 +- deepmd/kernels/cuda/dpa1/graph_compress.py | 4 +- .../kernels/cuda/dpa1/graph_energy_force.py | 10 ++-- deepmd/kernels/cuda/edge_force_virial.py | 10 ++-- deepmd/pt_expt/descriptor/dpa1.py | 2 +- deepmd/pt_expt/infer/deep_eval.py | 14 ++--- deepmd/pt_expt/model/edge_transform_output.py | 16 +++--- deepmd/pt_expt/model/ener_model.py | 18 +++--- deepmd/pt_expt/model/make_model.py | 28 +++++----- deepmd/pt_expt/train/training.py | 10 ++-- deepmd/pt_expt/utils/canonical_graph.py | 2 +- deepmd/pt_expt/utils/graph_csr.py | 2 +- deepmd/pt_expt/utils/serialization.py | 12 ++-- source/api_cc/include/DeepPotPTExpt.h | 6 +- source/api_cc/include/commonPT.h | 2 +- source/api_cc/src/DeepPotPTExpt.cc | 18 +++--- .../api_cc/tests/test_neighbor_list_data.cc | 2 +- source/op/pt/dpa1_graph_energy_force.cu | 12 ++-- source/op/pt/edge_force_virial.cu | 40 ++++++------- source/op/pt/graph_ops.h | 2 +- .../dpmodel/test_apply_pair_exclusion.py | 2 +- .../common/dpmodel/test_neighbor_graph.py | 6 +- .../dpmodel/test_neighbor_graph_builder.py | 4 +- .../pt_expt/descriptor/test_dpa1_cuda.py | 56 +++++++++---------- .../pt_expt/model/test_dpa1_graph_lower.py | 8 +-- .../tests/pt_expt/model/test_graph_export.py | 2 +- .../pt_expt/utils/test_canonical_graph.py | 2 +- source/tests/pt_expt/utils/test_graph_csr.py | 12 ++-- 31 files changed, 165 insertions(+), 165 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index a112619fce..f948714e4a 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -184,8 +184,8 @@ def from_dense_quartet( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( edge_index, edge_vec, @@ -200,8 +200,8 @@ def from_dense_quartet( edge_mask=edge_mask, destination_order=destination_order, destination_row_ptr=destination_row_ptr, - source_row_ptr=source_row_ptr, source_order=source_order, + source_row_ptr=source_row_ptr, destination_sorted=canonicalize, ) else: @@ -252,8 +252,8 @@ def from_dense_quartet( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( edge_index, edge_vec, @@ -268,8 +268,8 @@ def from_dense_quartet( edge_mask=edge_mask, destination_order=destination_order, destination_row_ptr=destination_row_ptr, - source_row_ptr=source_row_ptr, source_order=source_order, + source_row_ptr=source_row_ptr, destination_sorted=canonicalize, ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/csr.py b/deepmd/dpmodel/utils/neighbor_graph/csr.py index 8e0e8ea90c..d499a8f369 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/csr.py +++ b/deepmd/dpmodel/utils/neighbor_graph/csr.py @@ -65,11 +65,11 @@ def build_edge_csr( dtype as ``edge_index``. destination_row_ptr : Array Destination CSR offsets with shape ``(N + 1,)`` and dtype int64. - source_row_ptr : Array - Source 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 ------ @@ -126,8 +126,8 @@ def build_edge_csr( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) @@ -165,8 +165,8 @@ def attach_edge_csr( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( graph.edge_index, graph.edge_vec, @@ -181,8 +181,8 @@ def attach_edge_csr( edge_mask=edge_mask, destination_order=destination_order, destination_row_ptr=destination_row_ptr, - source_row_ptr=source_row_ptr, source_order=source_order, + source_row_ptr=source_row_ptr, destination_sorted=canonicalize, ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py index c3e5d37e64..d8c8f73c27 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py +++ b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py @@ -134,8 +134,8 @@ def neighbor_graph_from_ijs( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( edge_index, edge_vec, @@ -150,7 +150,7 @@ def neighbor_graph_from_ijs( edge_mask=edge_mask, destination_order=destination_order, destination_row_ptr=destination_row_ptr, - source_row_ptr=source_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 a7f7f52f4a..f702c2c6a1 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -61,10 +61,10 @@ class NeighborGraph: """(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_row_ptr: Array | None = field(default=None, kw_only=True) - """(N + 1,) int64 CSR offsets into ``source_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.""" @@ -333,8 +333,8 @@ def apply_pair_exclusion( edge_mask=xp.logical_and(graph.edge_mask, xp.astype(keep, xp.bool)), destination_order=None, destination_row_ptr=None, - source_row_ptr=None, source_order=None, + source_row_ptr=None, destination_sorted=False, ) if compact: diff --git a/deepmd/kernels/cuda/dpa1/graph_compress.py b/deepmd/kernels/cuda/dpa1/graph_compress.py index cddab24e6b..1bad0501d1 100644 --- a/deepmd/kernels/cuda/dpa1/graph_compress.py +++ b/deepmd/kernels/cuda/dpa1/graph_compress.py @@ -839,8 +839,8 @@ def dpa1_graph_compress_energy_force( if ( graph.destination_order is None or graph.destination_row_ptr is None - or graph.source_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" @@ -983,8 +983,8 @@ def dpa1_graph_compress_energy_force( graph.edge_mask, graph.destination_order, graph.destination_row_ptr, - graph.source_row_ptr, graph.source_order, + graph.source_row_ptr, graph.n_node, node_capacity, do_atomic_virial, diff --git a/deepmd/kernels/cuda/dpa1/graph_energy_force.py b/deepmd/kernels/cuda/dpa1/graph_energy_force.py index 6720a9d315..2cc262290e 100644 --- a/deepmd/kernels/cuda/dpa1/graph_energy_force.py +++ b/deepmd/kernels/cuda/dpa1/graph_energy_force.py @@ -61,8 +61,8 @@ def _fake( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_row_ptr: torch.Tensor, source_order: torch.Tensor, + source_row_ptr: torch.Tensor, atype: torch.Tensor, n_node: torch.Tensor, ownership: torch.Tensor, @@ -122,8 +122,8 @@ def _cpu( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_row_ptr: torch.Tensor, source_order: torch.Tensor, + source_row_ptr: torch.Tensor, atype: torch.Tensor, n_node: torch.Tensor, ownership: torch.Tensor, @@ -269,8 +269,8 @@ def _cpu( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, n_node, node_capacity, do_atomic_virial, @@ -369,8 +369,8 @@ def optional(t: torch.Tensor | None) -> torch.Tensor: if ( graph.destination_order is None or graph.destination_row_ptr is None - or graph.source_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" @@ -381,8 +381,8 @@ def optional(t: torch.Tensor | None) -> torch.Tensor: graph.edge_mask.contiguous(), graph.destination_order.contiguous(), graph.destination_row_ptr.contiguous(), - graph.source_row_ptr.contiguous(), graph.source_order.contiguous(), + graph.source_row_ptr.contiguous(), atype.contiguous(), graph.n_node, ownership.contiguous(), diff --git a/deepmd/kernels/cuda/edge_force_virial.py b/deepmd/kernels/cuda/edge_force_virial.py index 872044c950..22cbe58007 100644 --- a/deepmd/kernels/cuda/edge_force_virial.py +++ b/deepmd/kernels/cuda/edge_force_virial.py @@ -57,8 +57,8 @@ def _fake( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_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, @@ -97,8 +97,8 @@ def _cpu( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_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, @@ -168,8 +168,8 @@ def _canonical_cpu( device=edge_vec.device, ), destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, n_node_per_frame, node_capacity, want_atom_virial, @@ -209,8 +209,8 @@ def edge_force_virial( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_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, @@ -262,8 +262,8 @@ def edge_force_virial( edge_mask.contiguous(), destination_order.contiguous(), destination_row_ptr.contiguous(), - source_row_ptr.contiguous(), source_order.contiguous(), + source_row_ptr.contiguous(), n_node_per_frame, node_capacity, want_atom_virial, diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index 99c4630519..c1520402fb 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -66,8 +66,8 @@ def _has_graph_fields(graph: Any, fields: tuple[str, ...]) -> bool: _DESTINATION_CSR_FIELDS = ("destination_order", "destination_row_ptr") _DUAL_CSR_FIELDS = ( *_DESTINATION_CSR_FIELDS, - "source_row_ptr", "source_order", + "source_row_ptr", ) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 07f528d638..66efeddb89 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1767,7 +1767,7 @@ def _eval_model_graph( from the eval system at its exact (tight) edge count and feeds the positional schema ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, - destination_order, destination_row_ptr, source_row_ptr, source_order, + 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 @@ -1821,15 +1821,15 @@ def _eval_model_graph( 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, ) - source_order_t = torch.as_tensor( - graph.source_order, - device=DEVICE, - ) if self.metadata.get("lower_input_kind") == "dpa1_canonical": from deepmd.dpmodel.utils.neighbor_graph import ( @@ -1847,8 +1847,8 @@ def _eval_model_graph( n_local=n_node_t, destination_order=destination_order_t, destination_row_ptr=destination_row_ptr_t, - source_row_ptr=source_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) @@ -1876,8 +1876,8 @@ def _eval_model_graph( edge_mask_t, destination_order_t, destination_row_ptr_t, - source_row_ptr_t, source_order_t, + source_row_ptr_t, fparam_t, aparam_t, charge_spin_t, diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 4f9ff77c08..24c469c0a5 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -40,8 +40,8 @@ def edge_energy_deriv( n_node: torch.Tensor, destination_order: torch.Tensor | None = None, destination_row_ptr: torch.Tensor | None = None, - source_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, @@ -119,8 +119,8 @@ def edge_energy_deriv( and fused_scatter_available() and destination_order is not None and destination_row_ptr is not None - and source_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( @@ -130,16 +130,16 @@ def edge_energy_deriv( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, n_node, n_cap, do_atomic_virial, ) - return force, (atom_virial if do_atomic_virial else None), virial - force, atom_virial, virial = edge_force_virial( - g_e, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity - ) + 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 @@ -270,8 +270,8 @@ def fit_output_to_model_output_graph( n_node, graph.destination_order, graph.destination_row_ptr, - graph.source_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, diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 2710edb4c6..d0990ffe17 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -448,8 +448,8 @@ def forward_lower_graph_exportable( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_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, @@ -479,10 +479,10 @@ def forward_lower_graph_exportable( (E,) valid-edge mask (sample for tracing). destination_order (E,) destination-grouped edge permutation. - destination_row_ptr, source_row_ptr - (N + 1,) destination/source CSR offsets. 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. @@ -497,8 +497,8 @@ def forward_lower_graph_exportable( torch.nn.Module A traced module whose ``forward`` accepts ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, - destination_order, destination_row_ptr, source_row_ptr, - source_order, fparam, aparam, charge_spin)`` and returns a dict + 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 @@ -515,8 +515,8 @@ def forward_lower_graph_exportable( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, fparam=fparam, aparam=aparam, charge_spin=charge_spin, @@ -539,8 +539,8 @@ def fn( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_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, @@ -554,8 +554,8 @@ def fn( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, fparam, aparam, charge_spin, @@ -577,8 +577,8 @@ def fn( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, fparam, aparam, charge_spin, diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 2d40e6fd23..9885d599c6 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -351,8 +351,8 @@ def forward_common_lower_graph( edge_mask: torch.Tensor, destination_order: torch.Tensor | None = None, destination_row_ptr: torch.Tensor | None = None, - source_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, @@ -396,10 +396,10 @@ def forward_common_lower_graph( (E,) destination-grouped edge permutation. destination_row_ptr (N + 1,) destination CSR offsets into ``destination_order``. - source_row_ptr - (N + 1,) source CSR offsets into ``source_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. @@ -436,8 +436,8 @@ def forward_common_lower_graph( n_local=n_local, destination_order=destination_order, destination_row_ptr=destination_row_ptr, - source_row_ptr=source_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 @@ -633,8 +633,8 @@ def _call_common_graph( ng.edge_mask, ng.destination_order, ng.destination_row_ptr, - ng.source_row_ptr, ng.source_order, + ng.source_row_ptr, destination_sorted=ng.destination_sorted, do_atomic_virial=do_atomic_virial, fparam=fp, @@ -808,8 +808,8 @@ def forward_common_lower_graph_exportable( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_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, @@ -837,10 +837,10 @@ def forward_common_lower_graph_exportable( destination_order (E,) identity destination permutation. The exported ABI requires this canonical layout. - destination_row_ptr, source_row_ptr - (N + 1,) destination/source CSR offsets. 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. @@ -855,8 +855,8 @@ def forward_common_lower_graph_exportable( torch.nn.Module A traced module whose ``forward`` accepts ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, - destination_order, destination_row_ptr, source_row_ptr, - source_order, fparam, aparam, charge_spin)`` and returns a + 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``. """ @@ -865,8 +865,8 @@ def forward_common_lower_graph_exportable( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, atype.shape[0], destination_sorted=destination_sorted, ) @@ -881,8 +881,8 @@ def fn( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_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, @@ -899,8 +899,8 @@ def fn( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, destination_sorted=destination_sorted, do_atomic_virial=do_atomic_virial, fparam=fparam, @@ -917,8 +917,8 @@ def fn( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, fparam, aparam, charge_spin, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index da0a0087c7..08218835d7 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -710,8 +710,8 @@ def _trace_and_compile_graph( s_edge_mask, s_destination_order, s_destination_row_ptr, - s_source_row_ptr, s_source_order, + s_source_row_ptr, s_fparam, s_aparam, s_charge_spin, @@ -726,8 +726,8 @@ def fn( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_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, @@ -759,8 +759,8 @@ def fn( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, do_atomic_virial=False, fparam=fparam, aparam=aparam, @@ -799,8 +799,8 @@ def fn( s_edge_mask, s_destination_order, s_destination_row_ptr, - s_source_row_ptr, s_source_order, + s_source_row_ptr, s_fparam, s_aparam, s_charge_spin, @@ -1213,8 +1213,8 @@ def _forward_graph( ng.edge_mask, ng.destination_order, ng.destination_row_ptr, - ng.source_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 index ce449009d1..5db225ab1a 100644 --- a/deepmd/pt_expt/utils/canonical_graph.py +++ b/deepmd/pt_expt/utils/canonical_graph.py @@ -136,8 +136,8 @@ def canonical_graph_from_neighbor_graph( raise ValueError("canonical deployment requires destination-major payload") if ( graph.destination_row_ptr is None - or graph.source_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: diff --git a/deepmd/pt_expt/utils/graph_csr.py b/deepmd/pt_expt/utils/graph_csr.py index f001328e3c..1a0484931a 100644 --- a/deepmd/pt_expt/utils/graph_csr.py +++ b/deepmd/pt_expt/utils/graph_csr.py @@ -81,8 +81,8 @@ def validate_graph_csr_for_export( edge_mask: torch.Tensor, destination_order: torch.Tensor, destination_row_ptr: torch.Tensor, - source_row_ptr: torch.Tensor, source_order: torch.Tensor, + source_row_ptr: torch.Tensor, node_count: int, *, destination_sorted: bool, diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 0614d5527c..fe8357b858 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -392,7 +392,7 @@ def build_synthetic_graph_inputs( ``make_fx``. Inputs follow the positional order expected by ``forward_(common_)lower_graph``: ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, - destination_row_ptr, source_row_ptr, source_order, fparam, aparam, + destination_row_ptr, source_order, source_row_ptr, fparam, aparam, charge_spin)``. The system (``rng(42)``, ``box = rcut*3``, centered coords, ``atype[:, i] = @@ -490,8 +490,8 @@ def build_synthetic_graph_inputs( graph.edge_mask, graph.destination_order, graph.destination_row_ptr, - graph.source_row_ptr, graph.source_order, + graph.source_row_ptr, fparam, aparam, charge_spin, @@ -531,8 +531,8 @@ def build_synthetic_canonical_graph_inputs( edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, _fparam, _aparam, _charge_spin, @@ -545,8 +545,8 @@ def build_synthetic_canonical_graph_inputs( n_local=n_local, destination_order=destination_order, destination_row_ptr=destination_row_ptr, - source_row_ptr=source_row_ptr, source_order=source_order, + source_row_ptr=source_row_ptr, destination_sorted=True, ) compact = canonical_graph_from_neighbor_graph(graph) @@ -599,7 +599,7 @@ def _build_graph_dynamic_shapes( ---------- *sample_inputs : torch.Tensor | None ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, - destination_order, destination_row_ptr, source_row_ptr, source_order, + destination_order, destination_row_ptr, source_order, source_row_ptr, fparam, aparam, charge_spin)`` — 13 entries matching ``forward_lower_graph_exportable``. """ @@ -619,8 +619,8 @@ def _build_graph_dynamic_shapes( {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: n_node_total_dim + 1}, # source_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 diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 29e311022f..cd72c787e1 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -534,7 +534,7 @@ class DeepPotPTExpt : public DeepPotBackend { * * Positional AOTI input order matches the Python export ABI: * ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, - * destination_order, destination_row_ptr, source_row_ptr, source_order, + * destination_order, destination_row_ptr, source_order, source_row_ptr, * [fparam], [aparam], [charge_spin])``. * * @param[in] atype Per-node local-plus-halo types, shape ``(N,)`` int64. @@ -545,8 +545,8 @@ class DeepPotPTExpt : public DeepPotBackend { * @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_row_ptr Source 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, @@ -557,8 +557,8 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& edge_mask, const torch::Tensor& destination_order, const torch::Tensor& destination_row_ptr, - const torch::Tensor& source_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); diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 88293d83f3..12bd9754b5 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -377,8 +377,8 @@ struct GraphTensorPack { torch::Tensor edge_mask; torch::Tensor destination_order; torch::Tensor destination_row_ptr; - torch::Tensor source_row_ptr; torch::Tensor source_order; + torch::Tensor source_row_ptr; }; struct CanonicalGraphTensorPack { diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 45a18da7fb..8a81edd284 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -422,16 +422,16 @@ std::vector DeepPotPTExpt::run_model_graph( const torch::Tensor& edge_mask, const torch::Tensor& destination_order, const torch::Tensor& destination_row_ptr, - const torch::Tensor& source_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) { // 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_row_ptr, source_order}; + 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); } @@ -1008,7 +1008,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, 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_row_ptr, graph_pack.source_order, fparam_tensor, + graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, graph_aparam, charge_spin_tensor); } } else { @@ -1393,8 +1393,8 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, 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_row_ptr, - graph_tensors.source_order, fparam_tensor, aparam_tensor, + graph_tensors.destination_row_ptr, graph_tensors.source_order, + graph_tensors.source_row_ptr, fparam_tensor, aparam_tensor, charge_spin_tensor); } } else { @@ -2182,7 +2182,7 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, .findSchemaOrThrow("deepmd::build_graph_csr", "") .typed(); std::tie(graph_pack.destination_order, graph_pack.destination_row_ptr, - graph_pack.source_row_ptr, graph_pack.source_order) = + graph_pack.source_order, graph_pack.source_row_ptr) = build_graph_csr.call(edge_index, c10::SymInt(nnode), c10::SymInt(nedge)); } @@ -2192,7 +2192,7 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, 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_row_ptr, graph_pack.source_order, fparam_tensor, + 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(); diff --git a/source/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc index 37bbaa660b..ef83eae28b 100644 --- a/source/api_cc/tests/test_neighbor_list_data.cc +++ b/source/api_cc/tests/test_neighbor_list_data.cc @@ -150,8 +150,8 @@ TEST(TestNeighborListData, CompactCanonicalGraphDropsMaskedGuards) { 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_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); diff --git a/source/op/pt/dpa1_graph_energy_force.cu b/source/op/pt/dpa1_graph_energy_force.cu index c006b81daf..d1e9d56d3f 100644 --- a/source/op/pt/dpa1_graph_energy_force.cu +++ b/source/op/pt/dpa1_graph_energy_force.cu @@ -35,8 +35,8 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, torch::Tensor edge_mask, torch::Tensor destination_order, torch::Tensor destination_row_ptr, - torch::Tensor source_row_ptr, torch::Tensor source_order, + torch::Tensor source_row_ptr, torch::Tensor atype, torch::Tensor n_node, torch::Tensor ownership, @@ -129,10 +129,10 @@ dpa1_graph_energy_force(torch::Tensor edge_vec, // === Step 5. Scatter dE/d(edge_vec) into force / virial / atom virial. === // g_e and edge_vec_f are already in the compute precision; the per-node force // is a short neighbor sum and the per-frame virial reduces hierarchically. - auto fv = - edge_force_virial(g_e, edge_vec_f, edge_index, edge_mask, - destination_order, destination_row_ptr, source_row_ptr, - source_order, n_node, node_capacity, do_atomic_virial); + 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)}; } @@ -141,7 +141,7 @@ 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_row_ptr, Tensor source_order, Tensor atype, Tensor " + "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 " diff --git a/source/op/pt/edge_force_virial.cu b/source/op/pt/edge_force_virial.cu index 1fabd9d371..af0eb67861 100644 --- a/source/op/pt/edge_force_virial.cu +++ b/source/op/pt/edge_force_virial.cu @@ -57,8 +57,8 @@ __global__ void edge_force_virial_kernel( const bool* __restrict__ edge_mask, const index_t* __restrict__ destination_order, const long* __restrict__ destination_row_ptr, - const long* __restrict__ source_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; @@ -218,8 +218,8 @@ void launch_force_virial(long node_count, const torch::Tensor& edge_mask, const torch::Tensor& destination_order, const torch::Tensor& destination_row_ptr, - const torch::Tensor& source_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, @@ -235,9 +235,9 @@ void launch_force_virial(long node_count, edge_mask.numel() ? edge_mask.data_ptr() : nullptr, destination_order.numel() ? destination_order.data_ptr() : nullptr, - destination_row_ptr.data_ptr(), source_row_ptr.data_ptr(), - source_order.data_ptr(), force.data_ptr(), - node_virial.data_ptr()); + 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; @@ -261,8 +261,8 @@ std::tuple assemble_force_virial( const torch::Tensor& edge_mask, const torch::Tensor& destination_order, const torch::Tensor& destination_row_ptr, - const torch::Tensor& source_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); @@ -296,14 +296,14 @@ std::tuple assemble_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_row_ptr, - source_order, frame_row_ptr, force, node_virial, virial_partial, + 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_row_ptr, - source_order, frame_row_ptr, force, node_virial, virial_partial, + edge_mask, destination_order, destination_row_ptr, source_order, + source_row_ptr, frame_row_ptr, force, node_virial, virial_partial, virial, stream); } }); @@ -355,7 +355,7 @@ build_graph_csr(torch::Tensor edge_index, cursor.data_ptr(), source_order.data_ptr()); FORCE_CHECK_LAUNCH("build_graph_csr source order"); } - return {destination_order, destination_row_ptr, source_row_ptr, source_order}; + return {destination_order, destination_row_ptr, source_order, source_row_ptr}; } std::tuple edge_force_virial( @@ -365,16 +365,16 @@ std::tuple edge_force_virial( torch::Tensor edge_mask, torch::Tensor destination_order, torch::Tensor destination_row_ptr, - torch::Tensor source_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_row_ptr.is_cuda() && - source_order.is_cuda() && n_node_per_frame.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(), @@ -388,8 +388,8 @@ std::tuple edge_force_virial( "edge_force_virial: CSR row pointers must be int64"); TORCH_CHECK(destination_order.is_contiguous() && destination_row_ptr.is_contiguous() && - source_row_ptr.is_contiguous() && - source_order.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 || @@ -399,7 +399,7 @@ std::tuple edge_force_virial( "same int32 or int64 dtype"); return assemble_force_virial(node_count, edge_gradient, edge_vec, edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, n_node_per_frame, + source_order, source_row_ptr, n_node_per_frame, want_atom_virial); } @@ -442,7 +442,7 @@ canonical_edge_force_virial(torch::Tensor edge_gradient, 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_row_ptr, source_order, n_node_per_frame, + source_order, source_row_ptr, n_node_per_frame, want_atom_virial); } @@ -451,12 +451,12 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "build_graph_csr(Tensor edge_index, SymInt node_count, " "SymInt valid_edge_count) -> " "(Tensor destination_order, Tensor destination_row_ptr, " - "Tensor source_row_ptr, Tensor source_order)"); + "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_row_ptr, Tensor source_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)"); diff --git a/source/op/pt/graph_ops.h b/source/op/pt/graph_ops.h index 43c812e996..669513b2b5 100644 --- a/source/op/pt/graph_ops.h +++ b/source/op/pt/graph_ops.h @@ -123,8 +123,8 @@ std::tuple edge_force_virial( torch::Tensor edge_mask, torch::Tensor destination_order, torch::Tensor destination_row_ptr, - torch::Tensor source_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 98366b5794..9ff1dfb29b 100644 --- a/source/tests/common/dpmodel/test_apply_pair_exclusion.py +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion.py @@ -55,8 +55,8 @@ def test_exclusion_invalidates_derived_csr_views() -> None: assert out.destination_order is None assert out.destination_row_ptr is None - assert out.source_row_ptr is None assert out.source_order is None + assert out.source_row_ptr is None assert out.destination_sorted is False diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index dd4f695b53..b8b57c8b31 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -52,8 +52,8 @@ def test_preserves_payload_and_builds_stable_permutations(self) -> None: result_mask, destination_order, destination_row_ptr, - source_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) @@ -82,8 +82,8 @@ def test_canonicalizes_destination_payload_stably(self) -> None: result_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( edge_index, edge_vec, @@ -128,8 +128,8 @@ def test_canonicalization_rebuilds_an_untrusted_canonical_claim(self) -> None: 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_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, ) diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index c40febcbd4..4f45d22505 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -124,13 +124,13 @@ def test_canonicalization_is_explicit(self) -> None: self.assertFalse(generic.destination_sorted) self.assertIsNone(generic.destination_order) self.assertIsNone(generic.destination_row_ptr) - self.assertIsNone(generic.source_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_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, diff --git a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py index a346d9a44e..97920c65d8 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py @@ -289,12 +289,12 @@ def run(d, g, at, te): if graph.destination_row_ptr is not None else None ), - source_row_ptr=( - graph.source_row_ptr.cpu() if graph.source_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()) @@ -479,8 +479,8 @@ def run(d, g, at, te): edge_mask=graph.edge_mask.cpu(), destination_order=graph.destination_order.cpu(), destination_row_ptr=graph.destination_row_ptr.cpu(), - source_row_ptr=graph.source_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()) @@ -767,8 +767,8 @@ def test_unsorted_edge_stream_uses_permutation_csr(self) -> None: edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( graph.edge_index[:, permutation], graph.edge_vec[permutation], @@ -782,8 +782,8 @@ def test_unsorted_edge_stream_uses_permutation_csr(self) -> None: edge_mask=edge_mask, destination_order=destination_order, destination_row_ptr=destination_row_ptr, - source_row_ptr=source_row_ptr, source_order=source_order, + source_row_ptr=source_row_ptr, destination_sorted=False, ) type_embedding = des.type_embedding.call() @@ -853,8 +853,8 @@ def test_missing_csr_falls_back(self) -> None: graph, destination_order=None, destination_row_ptr=None, - source_row_ptr=None, source_order=None, + source_row_ptr=None, ) with _CudaLevel("0"): descriptor_ref, rotation_ref = des.call_graph( @@ -889,8 +889,8 @@ def test_empty_node_graph_has_zero_edge_gradient(self) -> None: 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_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) @@ -923,8 +923,8 @@ def test_empty_node_cpu_operator_has_zero_edge_gradient(self) -> None: edge_mask=torch.zeros(2, dtype=torch.bool), destination_order=order, destination_row_ptr=torch.zeros(1, dtype=torch.int64), - source_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) @@ -1185,8 +1185,8 @@ def test_parity_vs_separate_ops(self) -> None: graph.edge_mask, graph.destination_order, graph.destination_row_ptr, - graph.source_row_ptr, graph.source_order, + graph.source_row_ptr, graph.n_node, n, True, @@ -1218,8 +1218,8 @@ def test_missing_csr_declines_energy_force_fusion(self) -> None: graph, destination_order=None, destination_row_ptr=None, - source_row_ptr=None, source_order=None, + source_row_ptr=None, ) result = des.fused_energy_force_graph( fit, @@ -1258,8 +1258,8 @@ def test_level2_reuses_virtual_and_pair_exclusion_masks(self) -> None: graph.edge_mask, graph.destination_order, graph.destination_row_ptr, - graph.source_row_ptr, graph.source_order, + graph.source_row_ptr, ) with _CudaLevel("1"): @@ -1328,8 +1328,8 @@ def test_fused_energy_uses_owned_nodes_only(self) -> None: graph.edge_mask, graph.destination_order, graph.destination_row_ptr, - graph.source_row_ptr, graph.source_order, + graph.source_row_ptr, graph.n_node, n_node, True, @@ -1465,8 +1465,8 @@ def test_parity_vs_separate_ops(self) -> None: graph.edge_mask, graph.destination_order, graph.destination_row_ptr, - graph.source_row_ptr, graph.source_order, + graph.source_row_ptr, graph.n_node, n, True, @@ -1606,8 +1606,8 @@ def test_level2_permutation_csr_parity(self) -> None: edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( canonical_graph.edge_index[:, permutation], canonical_graph.edge_vec[permutation], @@ -1621,8 +1621,8 @@ def test_level2_permutation_csr_parity(self) -> None: edge_mask=edge_mask, destination_order=destination_order, destination_row_ptr=destination_row_ptr, - source_row_ptr=source_row_ptr, source_order=source_order, + source_row_ptr=source_row_ptr, destination_sorted=False, ) type_embedding = des.type_embedding.call() @@ -1668,7 +1668,7 @@ def test_linear_csr_builder(self) -> None: [0, 0, 1, 2, 2, 0, 0], dtype=torch.int64, device="cuda" ) edge_index = torch.stack([source, destination]) - destination_order, destination_row_ptr, source_row_ptr, source_order = ( + destination_order, destination_row_ptr, source_order, source_row_ptr = ( torch.ops.deepmd.build_graph_csr(edge_index, 3, 5) ) @@ -1722,8 +1722,8 @@ def _random_graph(self, device, n_frame=2, n_per_frame=9, n_edge=200): _topology_mask, dst_order, dst_row_ptr, - src_row_ptr, src_order, + src_row_ptr, ) = build_edge_csr( edge_index, torch.cat([g_e, edge_vec], dim=1), @@ -1739,8 +1739,8 @@ def _random_graph(self, device, n_frame=2, n_per_frame=9, n_edge=200): move(mask), move(dst_order), move(dst_row_ptr), - move(src_row_ptr), move(src_order), + move(src_row_ptr), move(n_node), total, ) @@ -1753,8 +1753,8 @@ def _reference( mask, dst_order, dst_row_ptr, - src_row_ptr, src_order, + src_row_ptr, n_node, total, ): @@ -1774,8 +1774,8 @@ def _fused( mask, dst_order, dst_row_ptr, - src_row_ptr, src_order, + src_row_ptr, n_node, total, ): @@ -1790,8 +1790,8 @@ def _fused( mask, dst_order, dst_row_ptr, - src_row_ptr, src_order, + src_row_ptr, n_node, total, True, @@ -1840,8 +1840,8 @@ def test_compact_canonical_parity(self) -> None: edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( edge_index, torch.cat((g_e, edge_vec), dim=1), @@ -1859,8 +1859,8 @@ def test_compact_canonical_parity(self) -> None: n_local=n_node, destination_order=destination_order, destination_row_ptr=destination_row_ptr, - source_row_ptr=source_row_ptr, source_order=source_order, + source_row_ptr=source_row_ptr, destination_sorted=True, ) compact = canonical_graph_from_neighbor_graph(graph) @@ -1872,8 +1872,8 @@ def test_compact_canonical_parity(self) -> None: edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, n_node, total, True, @@ -1923,8 +1923,8 @@ def test_many_small_frames(self) -> None: edge_mask, order, row_ptr, - row_ptr, order, + row_ptr, n_node_per_frame, frame_count, ) @@ -1980,8 +1980,8 @@ def test_ragged_many_frame_parity(self) -> None: _topology_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( edge_index, torch.cat([edge_gradient, edge_vec], dim=1), @@ -1996,8 +1996,8 @@ def test_ragged_many_frame_parity(self) -> None: edge_mask.cuda(), destination_order.cuda(), destination_row_ptr.cuda(), - source_row_ptr.cuda(), source_order.cuda(), + source_row_ptr.cuda(), n_node.cuda(), total, ) 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 4912525b9d..bd77cc1783 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -313,8 +313,8 @@ def test_graph_lower_symbolic_trace(self, attn_layer) -> None: em, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, fp, ap, cs, @@ -328,8 +328,8 @@ def test_graph_lower_symbolic_trace(self, attn_layer) -> None: em, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, fparam=fp, aparam=ap, do_atomic_virial=True, @@ -347,8 +347,8 @@ def test_graph_lower_symbolic_trace(self, attn_layer) -> None: em, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, fp, ap, cs, @@ -362,8 +362,8 @@ def test_graph_lower_symbolic_trace(self, attn_layer) -> None: em, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, fparam=fp, aparam=ap, do_atomic_virial=True, diff --git a/source/tests/pt_expt/model/test_graph_export.py b/source/tests/pt_expt/model/test_graph_export.py index 4b7e426ca2..1114b78848 100644 --- a/source/tests/pt_expt/model/test_graph_export.py +++ b/source/tests/pt_expt/model/test_graph_export.py @@ -79,8 +79,8 @@ def _graph_inputs(model): edge_mask, g.destination_order, g.destination_row_ptr, - g.source_row_ptr, g.source_order, + g.source_row_ptr, ) diff --git a/source/tests/pt_expt/utils/test_canonical_graph.py b/source/tests/pt_expt/utils/test_canonical_graph.py index 88a6cffc37..e7041d7f6e 100644 --- a/source/tests/pt_expt/utils/test_canonical_graph.py +++ b/source/tests/pt_expt/utils/test_canonical_graph.py @@ -30,8 +30,8 @@ def _generic_graph(physical_edge_count: int) -> NeighborGraph: n_local=torch.tensor([1], dtype=torch.int64), destination_order=torch.arange(storage_count, dtype=torch.int64), destination_row_ptr=row_ptr, - source_row_ptr=row_ptr.clone(), source_order=torch.arange(storage_count, dtype=torch.int64), + source_row_ptr=row_ptr.clone(), destination_sorted=True, ) diff --git a/source/tests/pt_expt/utils/test_graph_csr.py b/source/tests/pt_expt/utils/test_graph_csr.py index 57648489b4..be3eb66f13 100644 --- a/source/tests/pt_expt/utils/test_graph_csr.py +++ b/source/tests/pt_expt/utils/test_graph_csr.py @@ -25,8 +25,8 @@ def _csr_inputs(canonicalize: bool) -> tuple[torch.Tensor, ...]: edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = build_edge_csr( edge_index, edge_vec, @@ -39,8 +39,8 @@ def _csr_inputs(canonicalize: bool) -> tuple[torch.Tensor, ...]: edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) @@ -70,8 +70,8 @@ def test_rejects_nonidentity_canonical_order() -> None: edge_mask, destination_order, destination_row_ptr, - source_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) @@ -82,8 +82,8 @@ def test_rejects_nonidentity_canonical_order() -> None: edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, node_count=3, destination_sorted=True, ) @@ -95,8 +95,8 @@ def test_rejects_active_edge_outside_its_csr_row() -> None: edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, ) = _csr_inputs(canonicalize=True) destination_row_ptr = destination_row_ptr.clone() destination_row_ptr[1] = 1 @@ -107,8 +107,8 @@ def test_rejects_active_edge_outside_its_csr_row() -> None: edge_mask, destination_order, destination_row_ptr, - source_row_ptr, source_order, + source_row_ptr, node_count=3, destination_sorted=True, ) From 45bf817b9a95c93effdfa63406aa547783240b1e Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 13 Jul 2026 13:34:53 +0800 Subject: [PATCH 8/9] fix(lmp): expose device graph inference through C API Provide pointer-based C entry points and header-only wrappers so built-in and distributed Kokkos integrations share the supported C API path. --- doc/install/install-lammps.md | 12 +- source/api_c/include/c_api.h | 121 ++++++++++++++++- source/api_c/include/deepmd.hpp | 134 +++++++++++++++++- source/api_c/src/c_api.cc | 136 +++++++++++++++++++ source/api_c/tests/test_deepmd_exception.cc | 12 ++ source/api_c/tests/test_deeppot_universal.cc | 3 + source/api_cc/include/DeepPot.h | 7 +- source/api_cc/include/DeepPotPTExpt.h | 5 +- source/lmp/builtin.cmake | 25 +--- source/lmp/pair_deepmd_kokkos.cpp | 4 +- source/lmp/pair_deepmd_kokkos.h | 6 +- source/lmp/plugin/CMakeLists.txt | 9 ++ 12 files changed, 435 insertions(+), 39 deletions(-) diff --git a/doc/install/install-lammps.md b/doc/install/install-lammps.md index a61225ca5b..35278d6493 100644 --- a/doc/install/install-lammps.md +++ b/doc/install/install-lammps.md @@ -80,11 +80,11 @@ make -j4 make install ``` -To build the GPU-resident `pair_style deepmd/kk`, enable both the LAMMPS -Kokkos package and the DeePMD C++-API integration: +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 DEEPMD_LAMMPS_KOKKOS=ON \ +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 @@ -123,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/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 31b9e0626c..ec8b284717 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 1a68d0e227..62f7175ebb 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -940,7 +940,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 @@ -1175,6 +1201,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 1346bfbf20..21a0a86b56 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -1535,6 +1535,51 @@ const char* string_to_char(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) { @@ -1607,6 +1652,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 96f6942a65..700ab53a70 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" @@ -25,3 +26,14 @@ TEST(TestDeepmdException, deepmdexception_nofile) { ASSERT_THROW(deepmd::hpp::DeepPot("_no_such_file.pb"), 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); +} 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 5a9e83db34..262c653713 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -765,9 +765,10 @@ class DeepPot : public DeepBaseModel { * stored defaults. Empty vectors fall back to the stored default fparam and * to no aparam, so the two overloads coincide. * - * @param[in] fparam Runtime frame parameters, or empty for the model default. - * @param[in] aparam Runtime per-atom parameters (row-major [nloc, daparam]), - * or empty for none. + * @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. diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index cd72c787e1..b03ab5290e 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -322,8 +322,9 @@ class DeepPotPTExpt : public DeepPotBackend { * @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 Frame parameters, or empty for the stored default. - * @param[in] aparam Per-atom parameters, or empty for none. + * @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 diff --git a/source/lmp/builtin.cmake b/source/lmp/builtin.cmake index b4f05a2dd1..e051e5c24a 100644 --- a/source/lmp/builtin.cmake +++ b/source/lmp/builtin.cmake @@ -58,7 +58,6 @@ configure_file("${CMAKE_CURRENT_LIST_DIR}/deepmd_version.h.in" file(GLOB DEEPMD_LMP_SRC ${CMAKE_CURRENT_LIST_DIR}/*.cpp) find_package(DeePMD REQUIRED) -option(DEEPMD_LAMMPS_KOKKOS "Build the DeePMD Kokkos pair style" OFF) target_sources( lammps PRIVATE ${DEEPMD_LMP_SRC} @@ -69,29 +68,7 @@ target_sources( ${LAMMPS_SOURCE_DIR}/KSPACE/remap_wrap.cpp ${LAMMPS_SOURCE_DIR}/EXTRA-FIX/fix_ttm.cpp # for ttm ) -if(DEEPMD_LAMMPS_KOKKOS) - if(NOT PKG_KOKKOS) - message(FATAL_ERROR "DEEPMD_LAMMPS_KOKKOS requires PKG_KOKKOS") - endif() - if(NOT TARGET DeePMD::deepmd_cc) - message( - FATAL_ERROR - "DEEPMD_LAMMPS_KOKKOS requires the DeePMD C++ interface target") - endif() - target_link_libraries(lammps PUBLIC DeePMD::deepmd_cc) - target_compile_definitions(lammps PRIVATE DP_USE_CXX_API) - get_target_property(_deepmd_inc DeePMD::deepmd_cc - INTERFACE_INCLUDE_DIRECTORIES) - if(_deepmd_inc AND NOT _deepmd_inc MATCHES "-NOTFOUND$") - foreach(_inc IN LISTS _deepmd_inc) - if(NOT _inc MATCHES "^\\$<") - target_include_directories(lammps PUBLIC "${_inc}/deepmd") - endif() - endforeach() - endif() -else() - target_link_libraries(lammps PUBLIC DeePMD::deepmd_c) -endif() +target_link_libraries(lammps PUBLIC DeePMD::deepmd_c) target_include_directories( lammps PRIVATE ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_LIST_DIR} ${LAMMPS_SOURCE_DIR}/KSPACE ${LAMMPS_SOURCE_DIR}/EXTRA-FIX) diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index e5c3213913..9ecd68f055 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: LGPL-3.0-or-later -#if defined(LMP_KOKKOS) && defined(DP_USE_CXX_API) +#ifdef LMP_KOKKOS #include "pair_deepmd_kokkos.h" #include @@ -963,4 +963,4 @@ template class PairDeepMDKokkos; #endif } // namespace LAMMPS_NS -#endif // LMP_KOKKOS && DP_USE_CXX_API +#endif // LMP_KOKKOS diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index 54ab13e3ff..d09e9a7abc 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -1,6 +1,6 @@ // SPDX-License-Identifier: LGPL-3.0-or-later -// The device pair style requires both LAMMPS Kokkos and the DeePMD C++ API. -#if defined(LMP_KOKKOS) && defined(DP_USE_CXX_API) +// 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 @@ -152,4 +152,4 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { #endif #endif -#endif // LMP_KOKKOS && DP_USE_CXX_API +#endif // LMP_KOKKOS 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}) From e5d0dbefc4db0078f8254f216bef9e0c4530df4f Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 13 Jul 2026 13:38:45 +0800 Subject: [PATCH 9/9] fix(cuda): restore full architecture set on CUDA 13 Limit the CUDA 12.9 CCCL workaround to that release; CUDA 13 supports its current architecture set with -arch=all. --- source/op/pt/CMakeLists.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 4f07c19cc5..a6560e4fec 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -13,9 +13,9 @@ option( if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) find_package(CUDAToolkit REQUIRED) if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) - # CUDA 12.9 and 13.x CCCL fail to compile CUB/Thrust with -arch=all. + # 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 "14.0") + VERSION_LESS "13.0") set(CMAKE_CUDA_ARCHITECTURES all-major) else() set(CMAKE_CUDA_ARCHITECTURES all)