From 120fe0d6f9cf1e1f5804492d8a7477a4c4a12bdf Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 24 Oct 2026 23:15:21 +0800 Subject: [PATCH 01/15] feat(long-range): latent-charge long-range module for DPA4C A long-range module adds to the atomic energy of the energy fitting the interaction of latent charges that a charge network reads from the descriptor, E = E_SR + k_e/2 sum_{i != j} sum_c q_i^c q_j^c K(r_ij), with the LES kernel erf(alpha r)/r (the default) or a sum of Gaussians (SOG). Non-periodic systems take the direct sum over the pairs of a frame; periodic ones, with `periodic: true`, the sum over the images of the cell in reciprocal space, by the smooth particle-mesh Ewald method by default or over explicit wave numbers. A constraint shifts one channel so that it sums to the total charge of the charge state of the frame. Training (pt_expt): compiled training, DDP, and fine-tuning a pretrained model without long range into one with it, which takes the input statistics of the fine-tuning data for the added module. Frozen models: the graph lower reads the pair list and the mesh of the cell as topologies after its inputs; compressed DPA4C models frozen for a CUDA device evaluate the module in the compact canonical lower with the operator deepmd::latent_charge_sum (direct sum, SPME mesh, MPI exchange of the charges of a decomposed system). The C and C++ interfaces and LAMMPS (deepmd and deepmd/kk, single and multiple ranks, periodic and non-periodic) pass the node positions, the communicator and the cell through DP_DPA4ComputeCanonicalGraphGPU, which takes the cell of a periodic system after the communicator, and DP_DeepPotHasLongRange. Co-authored-by: yuzhibytedance --- .../dpmodel/atomic_model/base_atomic_model.py | 86 +- .../dpmodel/atomic_model/dp_atomic_model.py | 258 ++- .../dpmodel/atomic_model/inner_potential.py | 9 +- .../atomic_model/linear_atomic_model.py | 88 +- deepmd/dpmodel/descriptor/dpa4c.py | 42 +- .../dpmodel/descriptor/dpa4c_nn/__init__.py | 2 - .../dpmodel/descriptor/dpa4c_nn/geometry.py | 34 - deepmd/dpmodel/fitting/__init__.py | 4 + deepmd/dpmodel/fitting/latent_charge.py | 45 + deepmd/dpmodel/long_range/__init__.py | 34 + deepmd/dpmodel/long_range/base_long_range.py | 195 ++ deepmd/dpmodel/long_range/kernel.py | 419 ++++ deepmd/dpmodel/long_range/latent_charge.py | 730 ++++++ deepmd/dpmodel/long_range/reciprocal.py | 429 ++++ deepmd/dpmodel/model/dp_model.py | 2 + deepmd/dpmodel/model/make_model.py | 28 +- deepmd/dpmodel/model/model.py | 4 + deepmd/dpmodel/model/model_factory.py | 138 +- deepmd/dpmodel/model/native_spin_model.py | 4 + deepmd/dpmodel/model/spin_model.py | 5 + .../dpmodel/utils/neighbor_graph/__init__.py | 62 +- .../utils/neighbor_graph/derivatives.py | 68 + .../utils/neighbor_graph/pair_graph.py | 224 ++ .../utils/neighbor_graph/reciprocal_graph.py | 570 +++++ .../dpmodel/utils/neighbor_graph/topology.py | 271 +++ deepmd/jax/atomic_model/dp_atomic_model.py | 2 + deepmd/pd/model/model/__init__.py | 5 + deepmd/pt/model/model/__init__.py | 10 + deepmd/pt/utils/finetune.py | 4 +- deepmd/pt_expt/descriptor/dpa4c.py | 14 + deepmd/pt_expt/fitting/__init__.py | 4 + deepmd/pt_expt/fitting/latent_charge.py | 12 + deepmd/pt_expt/infer/deep_eval.py | 105 +- deepmd/pt_expt/kernels/dpa4c/canonical.py | 15 + .../kernels/dpa4c/canonical_long_range.py | 160 ++ deepmd/pt_expt/kernels/long_range.py | 475 ++++ deepmd/pt_expt/long_range/__init__.py | 18 + deepmd/pt_expt/long_range/base_long_range.py | 9 + deepmd/pt_expt/long_range/kernel.py | 55 + deepmd/pt_expt/long_range/latent_charge.py | 177 ++ deepmd/pt_expt/model/edge_transform_output.py | 155 +- deepmd/pt_expt/model/ener_model.py | 186 +- deepmd/pt_expt/model/get_model.py | 6 + deepmd/pt_expt/model/make_model.py | 197 +- deepmd/pt_expt/train/training.py | 104 +- deepmd/pt_expt/utils/serialization.py | 350 ++- deepmd/tf/model/model.py | 5 + deepmd/tf2/atomic_model/dp_atomic_model.py | 2 + deepmd/utils/argcheck.py | 305 +++ deepmd/utils/bridging.py | 3 + deepmd/utils/finetune.py | 12 + doc/model/index.rst | 1 + doc/model/long-range.md | 247 ++ source/api_c/include/c_api.h | 41 +- source/api_c/include/deepmd.hpp | 14 +- source/api_c/src/c_api.cc | 13 +- .../tests/test_deeppot_long_range_hpp.cc | 81 + source/api_cc/include/DeepPot.h | 39 +- source/api_cc/include/DeepPotPTExpt.h | 58 +- source/api_cc/src/DeepPot.cc | 27 +- source/api_cc/src/DeepPotPTExpt.cc | 268 ++- source/api_cc/src/longRangePTExpt.h | 160 ++ .../tests/dpa4c_canonical_test_common.h | 4 +- .../test_deeppot_dpa4c_long_range_ptexpt.cc | 603 +++++ source/install/test_cc_local.sh | 7 + source/lmp/pair_deepmd.cpp | 46 +- source/lmp/pair_deepmd.h | 3 + source/lmp/pair_deepmd_kokkos.cpp | 45 +- source/lmp/pair_deepmd_kokkos.h | 6 +- .../lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py | 9 +- .../tests/test_lammps_dpa4c_long_range_pt2.py | 439 ++++ source/op/pt/CMakeLists.txt | 14 +- source/op/pt/long_range/charge_sum.h | 68 + source/op/pt/long_range/direct_sum.cu | 979 ++++++++ source/op/pt/long_range/domain_exchange.cc | 167 ++ source/op/pt/long_range/domain_exchange.h | 44 + source/op/pt/long_range/frame_sums.cuh | 214 ++ source/op/pt/long_range/latent_charge_sum.cu | 235 ++ source/op/pt/long_range/mesh_sum.cu | 1375 +++++++++++ source/op/pt/long_range/ops.cc | 89 + .../dpmodel/test_linear_atomic_model.py | 3 + .../tests/common/dpmodel/test_long_range.py | 918 ++++++++ .../tests/common/dpmodel/test_zbl_bridging.py | 3 + source/tests/common/test_finetune_utils.py | 22 + source/tests/infer/gen_dpa4c_long_range.py | 264 +++ source/tests/pt_expt/descriptor/test_dpa4c.py | 16 + .../pt_expt/kernels/test_long_range_sum.py | 447 ++++ source/tests/pt_expt/model/test_long_range.py | 2002 +++++++++++++++++ source/tests/pt_expt/test_change_bias.py | 2 + source/tests/pt_expt/test_finetune.py | 73 + 90 files changed, 14971 insertions(+), 216 deletions(-) create mode 100644 deepmd/dpmodel/fitting/latent_charge.py create mode 100644 deepmd/dpmodel/long_range/__init__.py create mode 100644 deepmd/dpmodel/long_range/base_long_range.py create mode 100644 deepmd/dpmodel/long_range/kernel.py create mode 100644 deepmd/dpmodel/long_range/latent_charge.py create mode 100644 deepmd/dpmodel/long_range/reciprocal.py create mode 100644 deepmd/dpmodel/utils/neighbor_graph/pair_graph.py create mode 100644 deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py create mode 100644 deepmd/dpmodel/utils/neighbor_graph/topology.py create mode 100644 deepmd/pt_expt/fitting/latent_charge.py create mode 100644 deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py create mode 100644 deepmd/pt_expt/kernels/long_range.py create mode 100644 deepmd/pt_expt/long_range/__init__.py create mode 100644 deepmd/pt_expt/long_range/base_long_range.py create mode 100644 deepmd/pt_expt/long_range/kernel.py create mode 100644 deepmd/pt_expt/long_range/latent_charge.py create mode 100644 doc/model/long-range.md create mode 100644 source/api_c/tests/test_deeppot_long_range_hpp.cc create mode 100644 source/api_cc/src/longRangePTExpt.h create mode 100644 source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc create mode 100644 source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py create mode 100644 source/op/pt/long_range/charge_sum.h create mode 100644 source/op/pt/long_range/direct_sum.cu create mode 100644 source/op/pt/long_range/domain_exchange.cc create mode 100644 source/op/pt/long_range/domain_exchange.h create mode 100644 source/op/pt/long_range/frame_sums.cuh create mode 100644 source/op/pt/long_range/latent_charge_sum.cu create mode 100644 source/op/pt/long_range/mesh_sum.cu create mode 100644 source/op/pt/long_range/ops.cc create mode 100644 source/tests/common/dpmodel/test_long_range.py create mode 100644 source/tests/infer/gen_dpa4c_long_range.py create mode 100644 source/tests/pt_expt/kernels/test_long_range_sum.py create mode 100644 source/tests/pt_expt/model/test_long_range.py diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 720833b5cd..a330c6de79 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -222,6 +222,67 @@ def fused_decomposition(self) -> "tuple[Any, Any] | None": """ return None + def has_long_range(self) -> bool: + """Whether the atomic energy carries a long-range module beside a fitting. + + Generic capability (concrete default ``False``): a model wrapping the + atomic model asks it whether the energy holds a long-range interaction, + e.g. before combining it with a scheme defined for a learned local + energy alone. + """ + return False + + def required_topologies(self) -> tuple[str, ...]: + """Names of the topologies the graph forward reads beside the neighbor graph. + + Generic capability (concrete default: none) in the sense of + :meth:`uses_graph_lower`: a long-range module, whose interaction has no + cutoff, declares the topology it is evaluated on (see + :func:`~deepmd.dpmodel.utils.neighbor_graph.build_topologies`), and + every caller of the graph forward builds what is declared. + """ + return () + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of :meth:`required_topologies`. + + Generic capability (concrete default: none) in the sense of + :meth:`required_topologies`: a module whose topology depends on its + own state, such as the wave-number cutoff of a periodic long-range + module, which follows its kernel, gives it here, by topology name. + """ + return {} + + def build_topologies( + self, + coord: Any, + n_node: Any, + box: Any, + *, + atype: Any, + node_mask: Any = None, + pair_excl: Any = None, + ) -> dict[str, Any]: + """Build the topologies of :meth:`required_topologies` for one batch. + + See :func:`~deepmd.dpmodel.utils.neighbor_graph.build_topologies`; + the build parameters are those of :meth:`topology_params`. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_topologies, + ) + + return build_topologies( + self.required_topologies(), + coord, + n_node, + box, + node_mask, + atype=atype, + pair_excl=pair_excl, + topology_params=self.topology_params(), + ) + def supports_native_spin(self) -> bool: """Returns whether this atomic model consumes a per-atom spin input. @@ -494,6 +555,8 @@ def forward_common_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict: """Graph analogue of :meth:`forward_common_atomic` on the flat node axis. @@ -530,6 +593,13 @@ def forward_common_atomic_graph( (and, from there, the descriptor's ``call_graph``). ``None`` for non-parallel inference (default). Mirrors :meth:`forward_common_atomic`'s ``comm_dict`` on the dense route. + topologies + the topologies of :meth:`required_topologies` by name, built on the + node axis of ``graph``; ``None`` when none is required. + node_mask + flat mask of the nodes an enclosing model counts (real, owned and + not excluded there), combined with this model's own; ``None`` + imposes nothing further. Returns ------- @@ -538,6 +608,8 @@ def forward_common_atomic_graph( """ graph, atype_clamped, output_mask = self._prepare_graph_inputs(graph, atype) + if node_mask is not None: + output_mask = output_mask & node_mask ret_dict = self.forward_atomic_graph( graph, atype_clamped, @@ -546,6 +618,8 @@ def forward_common_atomic_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, + node_mask=output_mask, ) return self._finalize_atomic_ret(ret_dict, output_mask, atype) @@ -1014,9 +1088,18 @@ def model_forward( self.get_rcut(), pair_excl=self.pair_excl, ) + atype_flat = xp.reshape(atype, (-1,)) + topologies = self.build_topologies( + xp.reshape(coord, (-1, 3)), + graph.n_node, + box, + atype=atype_flat, + node_mask=atype_flat >= 0, + pair_excl=self.pair_excl, + ) atomic_ret = self.forward_common_atomic_graph( graph, - xp.reshape(atype, (-1,)), + atype_flat, fparam=fparam, aparam=( xp.reshape( @@ -1028,6 +1111,7 @@ def model_forward( ), charge_spin=charge_spin, spin=None if spin is None else xp.reshape(spin, (-1, 3)), + topologies=topologies, ) # The graph route works on a flat node axis; restore the # per-frame layout the dense route returns. diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index b1e0e98af9..b92deb413b 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -24,6 +24,10 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) +from deepmd.dpmodel.long_range import ( + BaseLongRange, + check_long_range_support, +) from deepmd.dpmodel.output_def import ( FittingOutputDef, ) @@ -102,6 +106,12 @@ class DPAtomicModel(BaseAtomicModel): :math:`\mathcal{F}` is the fitting network, and :math:`\mathbf{y}^i` is the predicted atomic property (energy, dipole, etc.). + An energy model may carry a long-range module beside its fitting: it reads + the same descriptor and adds its own atomic energy to the fitting output, + so the atomic energy :math:`E^i` is the complete potential, including the + interactions without a cutoff. The module is evaluated on the + NeighborGraph route. + Parameters ---------- descriptor @@ -111,6 +121,8 @@ class DPAtomicModel(BaseAtomicModel): type_map Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. + long_range + Long-range module added to the energy output of the fitting. """ @@ -121,12 +133,14 @@ def __init__( descriptor: BaseDescriptor, fitting: BaseFitting, type_map: list[str], + long_range: BaseLongRange | None = None, **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) self.descriptor = descriptor self.fitting_net = fitting self.fitting_net.reinit_exclude(self.atom_exclude_types) + self.long_range = long_range self.type_map = type_map self.add_chg_spin_ebd: bool = self.descriptor.has_chg_spin_ebd() # Structural capability: only descriptors with a native spin @@ -148,6 +162,8 @@ def __init__( self.add_spin_ebd: bool = ( self._supports_native_spin and self.descriptor.use_spin is not None ) + if self.long_range is not None: + self._check_long_range() super().init_out_stat() # === Reference of the isolated atoms === # The fitting references its output only when the preset fixes the @@ -155,6 +171,43 @@ def __init__( resolve_vacuum_ref(self.fitting_net, self.preset_out_bias) self.init_vacuum_conditions() + def _check_long_range(self) -> None: + """Validate that the long-range module can join this descriptor and fitting.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + ) + + check_long_range_support(self.descriptor) + if not self.descriptor.uses_graph_lower(): + raise NotImplementedError( + "a long-range module is evaluated on the NeighborGraph lower, " + "which this descriptor does not provide" + ) + if "energy" not in self.fitting_output_def().keys(): + raise ValueError( + "a long-range module adds to the energy output, which this " + "fitting does not have" + ) + if self.atom_exclude_types: + raise NotImplementedError( + "`atom_exclude_types` is not supported with a long-range module, " + "which attributes each pair energy to both atoms of the pair" + ) + if self.pair_exclude_types and any( + name != ALL_PAIRS for name in self.long_range.required_topologies() + ): + raise NotImplementedError( + "`pair_exclude_types` is not supported with a long-range module " + "that serves periodic systems, whose reciprocal-space sum is " + "global and excludes no pair" + ) + if self.long_range.uses_charge_spin() and not self.add_chg_spin_ebd: + raise ValueError( + "the long-range module reads the total charge from " + "`charge_spin`, but the descriptor is not configured with the " + "charge/spin condition (`add_chg_spin_ebd`)" + ) + def init_vacuum_conditions(self) -> None: """Build the conditioning inputs of the reference atoms for the type map. @@ -214,20 +267,49 @@ def uses_graph_lower(self) -> bool: return bool(self.descriptor.uses_graph_lower()) def fused_decomposition(self) -> "tuple[DPAtomicModel, None]": - """A descriptor-fitting model is its own learned part.""" + """A descriptor-fitting model is its own learned part. + + The learned part carries the long-range module of the model, if any; + a fused pipeline that does not evaluate long-range modules declines a + learned part with one. + """ return self, None + def has_long_range(self) -> bool: + """Whether this model carries a long-range module.""" + return self.long_range is not None + + def required_topologies(self) -> tuple[str, ...]: + """The topologies of the long-range module.""" + if self.long_range is None: + return () + return self.long_range.required_topologies() + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of the long-range module.""" + if self.long_range is None: + return {} + return self.long_range.topology_params() + def has_message_passing_across_ranks(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.has_message_passing_across_ranks()) def supports_edge_parallel(self) -> bool: - """Delegates to this model's own descriptor.""" - return bool(self.descriptor.supports_edge_parallel()) + """Delegates to this model's own descriptor. + + A long-range module sums over complete systems, so a model carrying + one does not run under domain decomposition. + """ + return self.long_range is None and bool( + self.descriptor.supports_edge_parallel() + ) def dense_lower_supports_comm(self) -> bool: - """Delegates to this model's own descriptor.""" - return bool(self.descriptor.dense_lower_supports_comm()) + """Delegates to this model's own descriptor; a long-range module has no dense lower.""" + return self.long_range is None and bool( + self.descriptor.dense_lower_supports_comm() + ) def uses_compact_edge_pairs(self) -> bool: """Delegates to this model's own descriptor.""" @@ -267,6 +349,8 @@ def set_case_embd(self, case_idx: int) -> None: typically concatenated with the output of the descriptor and fed into the fitting net. """ self.fitting_net.set_case_embd(case_idx) + if self.long_range is not None: + self.long_range.set_case_embd(case_idx) def mixed_types(self) -> bool: """If true, the model @@ -298,6 +382,12 @@ def enable_compression( ) -> None: """Call descriptor enable_compression(). + A long-range module keeps its networks: it reads the output of the + compressed descriptor like that of the uncompressed one. A compressed + descriptor folds the charge/spin condition into its tables, so a + constraint that reads the total charge from the condition is folded + to the same state, the default condition of the descriptor. + Parameters ---------- min_nbor_dist @@ -318,6 +408,8 @@ def enable_compression( table_stride_2, check_frequency, ) + if self.long_range is not None and self.long_range.uses_charge_spin(): + self.long_range.fold_charge_state(self.descriptor.get_default_chg_spin()) def forward_atomic( self, @@ -358,6 +450,11 @@ def forward_atomic( the result dict, defined by the `FittingOutputDef`. """ + if self.long_range is not None: + raise NotImplementedError( + "a long-range module is evaluated on the NeighborGraph lower " + "only; the dense (nlist) lower cannot evaluate this model" + ) nframes, nloc, nnei = nlist.shape atype = xp_take_first_n(extended_atype, 1, nloc) @@ -572,13 +669,16 @@ def forward_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict[str, Array]: """Graph analogue of :meth:`forward_atomic` on the flat node axis. Runs the descriptor ``call_graph`` then the fitting ``call_graph`` PER NODE and returns the raw fitting dict on the flat ``(N, *)`` axis (no reduction or masking; the wrapper handles those). ``fparam`` is gathered to nodes by - ``frame_id`` so each node sees its frame's parameter. + ``frame_id`` so each node sees its frame's parameter. The atomic energy + of the long-range module is added to the energy output of the fitting. Parameters ---------- @@ -604,6 +704,11 @@ def forward_atomic_graph( ``call_graph`` (the message-passing part). ``None`` for non-parallel inference (default). Mirrors :meth:`forward_atomic`'s ``comm_dict`` on the dense route. + topologies + the topologies of :meth:`required_topologies` by name. + node_mask + flat mask of the real and owned nodes, which carry the fields of + the long-range module. Returns ------- @@ -635,6 +740,22 @@ def forward_atomic_graph( graph.n_local, atype.shape[0], ) + # The long-range module reads the frame conditions of the real frames, + # with the defaults the descriptor and fitting fall back on. + n_node = graph.n_node + if self.long_range is not None: + # The constraint and a compressed descriptor serve one charge + # state; compressing the atomic model folds both to it. + if self.long_range.uses_charge_spin() and self.get_dim_chg_spin() == 0: + raise ValueError( + "the descriptor folds the charge/spin condition into its " + "compressed tables, while the long-range charge constraint " + "still reads it; compress the atomic model, which folds " + "both to one state" + ) + fparam, charge_spin_frame = self._frame_conditions( + fparam, charge_spin, n_node.shape[0], graph.edge_vec + ) # === Step 2. Vacuum reference nodes === # Opted-in descriptors evaluate the reference in the same graph call. # Legacy descriptors must retain the dense reference used by folding. @@ -683,7 +804,7 @@ def forward_atomic_graph( # its type, so along the closing of a pair the learned energy follows # the switch linearly while the analytical pair term takes over. gate_kwargs = {} if node_gate is None else {"node_gate": node_gate} - return self.fitting_net.call_graph( + ret = self.fitting_net.call_graph( gg, atype, gr=rot_mat, @@ -694,6 +815,71 @@ def forward_atomic_graph( **vacuum_kwargs, **gate_kwargs, ) + # === Step 5. Long-range module === + # The module reads the descriptor of the real nodes; the reference + # nodes appended above take no part in it. + if self.long_range is not None: + if topologies is None or node_mask is None: + raise ValueError( + "the long-range module of this model reads the topologies " + f"{self.required_topologies()} and the node mask, which were " + "not given" + ) + ret["energy"] = ret["energy"] + self.long_range.call_graph( + gg, + atype, + node_mask, + n_node, + topologies, + fparam=fparam, + aparam=aparam_node, + charge_spin=charge_spin_frame, + ) + return ret + + def _frame_conditions( + self, + fparam: Array | None, + charge_spin: Array | None, + nf: int, + reference: Array, + ) -> tuple[Array | None, Array | None]: + """Frame parameters and charge/spin condition with their defaults resolved. + + Parameters + ---------- + fparam + Frame parameters with shape (nf, numb_fparam), or None. + charge_spin + Charge/spin condition with shape (nf, 2) or (1, 2), or None. + nf + The number of frames. + reference + An array giving the namespace, dtype and device of the defaults. + + Returns + ------- + tuple + The frame parameters with shape (nf, numb_fparam) and the + charge/spin condition with shape (nf, 2), each None when the model + takes no such input or it has no default. + """ + xp = array_api_compat.array_namespace(reference) + device = array_api_compat.device(reference) + if fparam is None and self.has_default_fparam(): + default_fparam = xp.asarray( + self.get_default_fparam(), dtype=reference.dtype, device=device + ) + fparam = xp.tile(xp.reshape(default_fparam, (1, -1)), (nf, 1)) + if not self.add_chg_spin_ebd: + return fparam, None + if charge_spin is None: + default_cs = self.descriptor.get_default_chg_spin() + if default_cs is None: + return fparam, None + charge_spin = xp.asarray(default_cs, dtype=reference.dtype, device=device) + charge_spin = xp.broadcast_to(xp.reshape(charge_spin, (-1, 2)), (nf, 2)) + return fparam, charge_spin def compute_or_load_stat( self, @@ -724,6 +910,10 @@ def compute_or_load_stat( self.fitting_net.compute_input_stats( wrapped_sampler, stat_file_path=stat_file_path ) + if self.long_range is not None: + self.long_range.compute_input_stats( + wrapped_sampler, stat_file_path=stat_file_path + ) self._collect_and_set_observed_type( wrapped_sampler, stat_file_path, preset_observed_type ) @@ -747,6 +937,8 @@ def change_type_map( else None, ) self.fitting_net.change_type_map(type_map=type_map) + if self.long_range is not None: + self.long_range.change_type_map(type_map=type_map) self.init_vacuum_conditions() def compute_fitting_input_stat( @@ -756,6 +948,9 @@ def compute_fitting_input_stat( ) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + The networks of the long-range module take the inputs of the fitting + and receive their statistics as well. + Parameters ---------- sample_merged : Union[Callable[[], list[dict]], list[dict]] @@ -772,6 +967,37 @@ def compute_fitting_input_stat( protection=self.data_stat_protect, stat_file_path=stat_file_path, ) + if self.long_range is not None: + self.long_range.compute_input_stats( + sample_merged, + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) + + def compute_long_range_input_stat( + self, + sampled_func: Callable[[], list[dict]], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics of the long-range module alone. + + A fine-tuning run that adds a long-range module to a pretrained model + calls it for the new module, whose statistics, unlike those of the + inherited descriptor and fitting, the pretrained model cannot provide. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data + systems. + stat_file_path + The path to the stat file. + """ + self.long_range.compute_input_stats( + self._make_wrapped_sampler(sampled_func), + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) def serialize(self) -> dict: dd = super().serialize() @@ -785,6 +1011,12 @@ def serialize(self) -> dict: "fitting": self.fitting_net.serialize(), } ) + # A record with a long-range module is version 3, which a reader that + # cannot evaluate the module rejects; every other record keeps + # version 2. + if self.long_range is not None: + dd["@version"] = 3 + dd["long_range"] = self.long_range.serialize() return dd # for subclass overridden @@ -792,17 +1024,27 @@ def serialize(self) -> dict: """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls: type | None = BaseLongRange + """The base long-range module class; ``None`` for a backend without one.""" @classmethod def deserialize(cls, data: dict[str, Any]) -> "DPAtomicModel": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 2, 2) + check_version_compatibility(data.pop("@version", 1), 3, 2) data.pop("@class") data.pop("type") descriptor_obj = cls.base_descriptor_cls.deserialize(data.pop("descriptor")) fitting_obj = cls.base_fitting_cls.deserialize(data.pop("fitting")) data["descriptor"] = descriptor_obj data["fitting"] = fitting_obj + long_range = data.pop("long_range", None) + if long_range is not None: + if cls.base_long_range_cls is None: + raise NotImplementedError( + "this backend cannot evaluate the long-range module of the " + "model record" + ) + data["long_range"] = cls.base_long_range_cls.deserialize(long_range) obj = super().deserialize(data) return obj diff --git a/deepmd/dpmodel/atomic_model/inner_potential.py b/deepmd/dpmodel/atomic_model/inner_potential.py index e9c96c96ea..749e65ee14 100644 --- a/deepmd/dpmodel/atomic_model/inner_potential.py +++ b/deepmd/dpmodel/atomic_model/inner_potential.py @@ -607,12 +607,15 @@ def forward_atomic_graph( charge_spin: Any = None, spin: Any = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Any = None, ) -> dict: """Evaluate the analytical per-atom energy on the flat node axis. - ``fparam``/``aparam``/``charge_spin``/``spin``/``comm_dict`` are - accepted for pipeline-signature compatibility and ignored (the - analytical term conditions on geometry and types only). + ``fparam``/``aparam``/``charge_spin``/``spin``/``comm_dict``/ + ``topologies``/``node_mask`` are accepted for pipeline-signature + compatibility and ignored (the analytical term conditions on + geometry and types only; the caller masks its outputs). Parameters ---------- diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index ab4939a30e..e0476272a3 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -12,6 +12,9 @@ from deepmd.dpmodel.array_api import ( Array, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, +) from deepmd.dpmodel.utils.nlist import ( build_multiple_neighbor_list, get_multiple_nlist_key, @@ -34,6 +37,9 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from .inner_potential import ( + InnerPotentialAtomicModel, +) from .pairtab_atomic_model import ( PairTabAtomicModel, ) @@ -112,6 +118,40 @@ def __init__( + ", ".join(f"{type(m).__name__}={dim_of(m)}" for m in models) ) + # A long-range module attributes each pair energy to both atoms of + # the pair. It is not combined with a bridging term, which defines the + # close-pair limit of a learned local energy, nor with atom exclusion, + # which removes atomic outputs. + if any(isinstance(m, InnerPotentialAtomicModel) for m in models) and any( + m.has_long_range() for m in models + ): + raise NotImplementedError( + "an inner_potential bridging term cannot be combined with a " + "sub-model carrying a long-range module" + ) + if self.atom_exclude_types and any(m.has_long_range() for m in models): + raise NotImplementedError( + "`atom_exclude_types` is not supported with a sub-model carrying " + "a long-range module" + ) + # The composition builds every topology once for all sub-models: the + # pairs of a periodic batch are guards only, which would silence a + # module serving non-periodic systems alone, and a reciprocal-space + # topology is built with the parameters of a single module. + declared = [m.required_topologies() for m in models if m.has_long_range()] + periodic = any(name != ALL_PAIRS for names in declared for name in names) + if len(declared) > 1 and periodic: + raise NotImplementedError( + "a sub-model whose long-range module serves periodic systems " + "takes no other sub-model carrying a long-range module beside it" + ) + if self.pair_exclude_types and periodic: + raise NotImplementedError( + "`pair_exclude_types` is not supported with a sub-model whose " + "long-range module serves periodic systems, whose " + "reciprocal-space sum is global and excludes no pair" + ) + self.models = models self.type_map = type_map self._rebuild_mapping_state() @@ -377,6 +417,8 @@ def compute_fitting_input_stats( for model in self.models: if isinstance(model, DPAtomicModel): model.fitting_net.compute_input_stats(sampled_func) + if model.long_range is not None: + model.long_range.compute_input_stats(sampled_func) def supports_native_spin(self) -> bool: """Spin-capable when ANY child consumes the spin input. @@ -389,6 +431,30 @@ def supports_native_spin(self) -> bool: """ return any(m.supports_native_spin() for m in self.models) + def has_long_range(self) -> bool: + """ANY child carrying a long-range module makes the composition carry one.""" + return any(m.has_long_range() for m in self.models) + + def required_topologies(self) -> tuple[str, ...]: + """Every topology a child reads, each named once.""" + return tuple( + dict.fromkeys( + name for model in self.models for name in model.required_topologies() + ) + ) + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of the sub-models. + + At most one sub-model declares a reciprocal-space topology, the only + kind with build parameters, so that the parameters of every topology + come from one sub-model. + """ + params: dict[str, dict[str, Any]] = {} + for model in self.models: + params.update(model.topology_params()) + return params + def forward_atomic_graph( self, graph: Any, @@ -398,14 +464,18 @@ def forward_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict[str, Array]: """Graph-route linear combination on the flat node axis. - Every child consumes the SAME graph, so on autograd backends the - shared ``graph.edge_vec`` leaf makes the summed energy's force and - virial exactly the sum of the children's -- one edge backward - covers the whole composition (this is what makes analytical - bridging terms compose with the learned model for free). + Every child consumes the SAME graph and topologies, so on autograd + backends the shared geometry leaves make the summed energy's force + and virial exactly the sum of the children's -- one backward covers + the whole composition (this is what makes analytical bridging terms + compose with the learned model for free). Each child evaluates its + complete energy, including its long-range module, before its weight + applies. Only constant weights are supported here (``"sum"``/``"mean"`` or a per-child float list); the distance-switched ZBL-interpolation @@ -429,6 +499,12 @@ def forward_atomic_graph( flat (N, 3) per-node spin, forwarded to every child. comm_dict MPI communication metadata, forwarded to every child. + topologies + the topologies of :meth:`required_topologies`, forwarded to every + child. + node_mask + flat mask of the nodes whose outputs count, forwarded to every + child, which combines it with its own exclusions. Returns ------- @@ -472,6 +548,8 @@ def forward_atomic_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, + node_mask=node_mask, ) contrib = ret["energy"] * ww energy = contrib if energy is None else energy + contrib diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index e7398bd3c3..d43c9255bd 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -90,7 +90,6 @@ OrderedPairFiLM, SpinChannels, build_angular_basis, - build_moment_indices, canonicalize_charge_spin, degree_offsets, derive_bispectrum_ranks, @@ -467,12 +466,9 @@ def __init__( # === Step 4. Lay out the flat moment payload === # Degree zero owns the leading `channels` entries of the flat layout, - # so the non-scalar block is exactly its complement and needs no - # separate degree index. The spin families, when present, are appended - # after the geometric degrees, leaving every geometric offset intact. - channel_index, harmonic_index = build_moment_indices(self.degree_channels) - self.angular_channel_index = channel_index[self.channels :] - self.angular_harmonic_index = harmonic_index[self.channels :] + # and every higher degree follows in its own harmonic-major block. The + # spin families, when present, are appended after the geometric + # degrees, leaving every geometric offset intact. self.degree_offsets = degree_offsets(self.degree_channels) # === Step 5. Initialize the output calibration state === @@ -1152,30 +1148,30 @@ def aggregate_moments( ) xp = array_api_compat.array_namespace(amplitude) - device = array_api_compat.device(amplitude) - channel_index = xp_asarray_nodetach( - xp, - self.angular_channel_index, - device=device, - ) - harmonic_index = xp_asarray_nodetach( - xp, - self.angular_harmonic_index, - device=device, - ) + n_edge = amplitude.shape[0] # Payload layout: [chi^2, chi^4, degree zero (C_0), degrees one and # above (S - C_0)]. Degree zero is the whole edge amplitude under its - # single envelope, while the non-scalar block gathers an amplitude and - # a harmonic per flat moment coordinate and carries a second envelope. + # single envelope. Degree l is the outer product of its 2l + 1 + # harmonics with the leading C_l channels of the amplitude, laid out + # harmonic-major and carrying a second envelope. The products are + # formed by broadcasting, so their backward reduces over the + # broadcast axes. + angular = [ + xp.reshape( + basis[:, degree * degree : (degree + 1) * (degree + 1), None] + * amplitude[:, None, :width], + (n_edge, (2 * degree + 1) * width), + ) + for degree, width in enumerate(self.degree_channels) + if degree > 0 + ] envelope_squared = envelope * envelope parts = [ envelope_squared[:, None], (envelope_squared * envelope_squared)[:, None], amplitude, - xp.take(amplitude, channel_index, axis=1) - * xp.take(basis, harmonic_index, axis=1) - * envelope[:, None], + xp.concat(angular, axis=1) * envelope[:, None], ] if spin_payload is not None: parts.append(spin_payload) diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py index 460c7970e6..8bd43d4cc0 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py @@ -14,7 +14,6 @@ from .geometry import ( MAX_ANGULAR_DEGREE, build_angular_basis, - build_moment_indices, degree_offsets, derive_degree_channels, packed_l2_to_stf, @@ -41,7 +40,6 @@ "SpinChannels", "build_angular_basis", "build_bispectrum_layout", - "build_moment_indices", "canonicalize_charge_spin", "degree_offsets", "derive_bispectrum_ranks", diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py index 766a9d060c..a9378f4ea1 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py @@ -11,7 +11,6 @@ ) import array_api_compat -import numpy as np if TYPE_CHECKING: from collections.abc import ( @@ -236,36 +235,3 @@ def degree_offsets(degree_channels: Sequence[int]) -> tuple[int, ...]: for degree, width in enumerate(degree_channels): offsets.append(offsets[-1] + (2 * degree + 1) * int(width)) return tuple(offsets) - - -def build_moment_indices( - degree_channels: Sequence[int], -) -> tuple[np.ndarray, np.ndarray]: - """Build channel and harmonic indices for one flat edge payload. - - Every degree reads the leading channels of the shared radial map, so the - channel index of a degree is simply ``range(degree_channels[degree])``. - - Parameters - ---------- - degree_channels - Number of channels for degrees zero through ``lmax``. - - Returns - ------- - channel_index - Edge-amplitude channel indices with shape ``(S,)``. - harmonic_index - Packed harmonic indices with shape ``(S,)``, where - ``S = sum((2 * l + 1) * degree_channels[l])``. - """ - channel_index = [] - harmonic_index = [] - for degree, width in enumerate(degree_channels): - for component in range(2 * degree + 1): - channel_index.extend(range(int(width))) - harmonic_index.extend([degree * degree + component] * int(width)) - return ( - np.asarray(channel_index, dtype=np.int64), - np.asarray(harmonic_index, dtype=np.int64), - ) diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 37b0a1b731..126b51e721 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -14,6 +14,9 @@ from .invar_fitting import ( InvarFitting, ) +from .latent_charge import ( + LatentChargeFitting, +) from .make_base_fitting import ( make_base_fitting, ) @@ -29,6 +32,7 @@ "DipoleFitting", "EnergyFittingNet", "InvarFitting", + "LatentChargeFitting", "PolarFitting", "PropertyFittingNet", "SeZMEnergyFittingNet", diff --git a/deepmd/dpmodel/fitting/latent_charge.py b/deepmd/dpmodel/fitting/latent_charge.py new file mode 100644 index 0000000000..4967879f1f --- /dev/null +++ b/deepmd/dpmodel/fitting/latent_charge.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) + +from .invar_fitting import ( + InvarFitting, +) + + +@fitting_check_output +class LatentChargeFitting(InvarFitting): + r"""Per-atom latent charges predicted from the invariant descriptor. + + The network maps the descriptor of an atom, together with the frame + parameters, atomic parameters and case embedding of the energy fitting it + accompanies, to ``dim_out`` latent charge channels in units of the + elementary charge :math:`e`. The charges are an intermediate field of a + long-range module rather than a model output: the module constrains and + pairs them, and the potential energy is the only quantity they are + trained through. + + The constructor arguments are those of :class:`InvarFitting` with the + output variable fixed to ``latent_charge``. + """ + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def serialize(self) -> dict: + data = super().serialize() + data["type"] = "latent_charge" + return data diff --git a/deepmd/dpmodel/long_range/__init__.py b/deepmd/dpmodel/long_range/__init__.py new file mode 100644 index 0000000000..4a4aad0d7c --- /dev/null +++ b/deepmd/dpmodel/long_range/__init__.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Long-range modules of the atomic energy of a descriptor-fitting model. + +A long-range module is a component of the atomic energy model beside its +energy fitting: it reads the descriptor features shared with that fitting and +the topologies it declares, and contributes an energy per atom. Every consumer +of the atomic model therefore evaluates the complete potential. +""" + +from .base_long_range import ( + BaseLongRange, + check_long_range_support, + make_base_long_range, +) +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, + LESKernel, + SOGKernel, +) +from .latent_charge import ( + LatentChargeLongRange, +) + +__all__ = [ + "COULOMB_CONSTANT", + "BaseInteractionKernel", + "BaseLongRange", + "LESKernel", + "LatentChargeLongRange", + "SOGKernel", + "check_long_range_support", + "make_base_long_range", +] diff --git a/deepmd/dpmodel/long_range/base_long_range.py b/deepmd/dpmodel/long_range/base_long_range.py new file mode 100644 index 0000000000..d00617fc1c --- /dev/null +++ b/deepmd/dpmodel/long_range/base_long_range.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + ABC, + abstractmethod, +) +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.plugin import ( + PluginVariant, + make_plugin_registry, +) + + +def check_long_range_support(descriptor: Any) -> None: + """Admit only the descriptors whose combination with long-range modules is validated. + + A long-range module reads nothing but the invariant node features every + graph descriptor returns, so this gate is a statement of validation, not + of implementation: a further descriptor is admitted by adding it here once + its combination has been tested. + + Parameters + ---------- + descriptor + The descriptor of the atomic model receiving the long-range module. + + Raises + ------ + NotImplementedError + If the descriptor is not among the validated ones. + """ + from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, + ) + + if not isinstance(descriptor, DescrptDPA4C): + raise NotImplementedError( + "long-range modules are currently validated with the DPA4C " + f"descriptor only, got {type(descriptor).__name__}" + ) + + +def make_base_long_range(t_tensor: Any) -> type: + """Make the base class of the long-range modules of one backend. + + Each backend owns its registry, so a configuration builds the long-range + module native to the backend that reads it. + + Parameters + ---------- + t_tensor + The type of the tensor, used in the type hints. + """ + + class BLR(ABC, PluginVariant, make_plugin_registry("long range")): + """Long-range contribution to the atomic energy of a descriptor-fitting model. + + A long-range module reads the invariant node features of the + descriptor and the topologies it declares, and returns an energy per + node in eV. The atomic model adds it to the output of its energy + fitting before the output statistics and masks, so every consumer of + the atomic model sees one complete potential. The module owns the + attribution of its energy to the nodes; the per-frame energy is their + sum. + """ + + @abstractmethod + def required_topologies(self) -> tuple[str, ...]: + """Names of the topologies the module reads beside the neighbor graph.""" + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies, by name; none by default.""" + return {} + + @abstractmethod + def uses_charge_spin(self) -> bool: + """Whether the module reads the frame-level charge/spin condition.""" + + @abstractmethod + def call_graph( + self, + descriptor: t_tensor, + atype: t_tensor, + node_mask: t_tensor, + n_node: t_tensor, + topologies: dict[str, Any], + fparam: t_tensor | None = None, + aparam: t_tensor | None = None, + charge_spin: t_tensor | None = None, + ) -> t_tensor: + """Evaluate the atomic energy of the module on the flat node axis. + + Parameters + ---------- + descriptor + Invariant node features with shape ``(N, dim_descrpt)``. + atype + Node types with shape ``(N,)``; invalid nodes carry any valid + type. + node_mask + Nodes carrying fields with shape ``(N,)``: the real and owned + nodes. + n_node + Per-frame node counts with shape ``(nf,)``. + topologies + The topologies of :meth:`required_topologies` by name. + fparam + Frame parameters with shape ``(nf, numb_fparam)``, defaults + resolved. + aparam + Atomic parameters with shape ``(N, numb_aparam)``. + charge_spin + Frame charge/spin condition with shape ``(nf, 2)``, defaults + resolved, when the model takes one. + + Returns + ------- + t_tensor + The atomic energy with shape ``(N, 1)``, in eV. + """ + + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics of the networks of the module.""" + + def change_type_map(self, type_map: list[str]) -> None: + """Remap the type-dependent state of the module to a new type map.""" + + def set_case_embd(self, case_idx: int) -> None: + """Select the case embedding of the networks of the module.""" + + def fold_charge_state(self, charge_spin: list[float]) -> None: + """Specialize the module to one frame charge/spin condition. + + Called when the descriptor of the model is compressed, which folds + the condition into its tables, for a module whose + :meth:`uses_charge_spin` holds. + + Raises + ------ + NotImplementedError + If the module cannot be specialized. + """ + raise NotImplementedError( + f"{type(self).__name__} cannot be specialized to a charge/spin " + "condition, which the compression of the model requires" + ) + + def charge_state_constants(self) -> list[t_tensor]: + """The tensors that carry the folded charge/spin condition. + + A deployment that serves another condition writes the outputs of + :meth:`rebuild_charge_state` over them, in this order. + """ + return [] + + def rebuild_charge_state(self, charge_spin: t_tensor) -> list[t_tensor]: + """Build the tensors of :meth:`charge_state_constants` for a condition. + + Parameters + ---------- + charge_spin + The condition with shape ``(1, 2)``. + """ + return [] + + @abstractmethod + def serialize(self) -> dict: + """Serialize the module to a dict.""" + + @classmethod + def deserialize(cls, data: dict) -> "BLR": + """Deserialize a module, dispatching on its registered type.""" + if cls is BLR: + return BLR.get_class_by_type(data["type"]).deserialize(data) + raise NotImplementedError(f"Not implemented in class {cls.__name__}") + + return BLR + + +BaseLongRange = make_base_long_range(np.ndarray) diff --git a/deepmd/dpmodel/long_range/kernel.py b/deepmd/dpmodel/long_range/kernel.py new file mode 100644 index 0000000000..911dfba5dd --- /dev/null +++ b/deepmd/dpmodel/long_range/kernel.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Pair interaction kernels of the latent-charge long-range modules. + +A kernel is a function :math:`K(r)` of the distance between two charges, in +Å⁻¹, so that charges in units of :math:`e` interact with the energy +:math:`k_e q_i q_j K(r_{ij})` in eV, where :math:`k_e` is the Coulomb +constant. A kernel owns its parameters; widths are stored through their +logarithm, which keeps them positive however training moves them. +""" + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, + xp_asarray_nodetach, + xp_erf, +) +from deepmd.dpmodel.common import ( + PRECISION_DICT, + NativeOP, + to_numpy_array, +) +from deepmd.utils.plugin import ( + make_plugin_registry, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +COULOMB_CONSTANT = 14.3996454784255 +r"""Coulomb constant :math:`e^2 / (4 \pi \varepsilon_0)` in eV Å.""" + + +class BaseInteractionKernel(NativeOP, make_plugin_registry("interaction kernel")): + """Pair interaction kernel :math:`K(r)` in Å⁻¹ with its own parameters. + + Parameters + ---------- + trainable + Whether the parameters of the kernel are trained. + precision + The precision in which the parameters are stored; the kernel is + evaluated in the precision of its input. + """ + + PARAMETER_NAMES: tuple[str, ...] = () + """Attributes holding the trainable parameters of the kernel.""" + + def __init__(self, trainable: bool = True, precision: str = "default") -> None: + self.trainable = bool(trainable) + self.precision = precision + self.prec = PRECISION_DICT[precision.lower()] + + def call(self, r_sq: Array) -> Array: + """Evaluate the kernel on squared pair distances. + + Parameters + ---------- + r_sq + Squared distances with shape ``(P,)``, in Ų. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + return self.evaluate(r_sq, self.coefficients(r_sq)) + + def coefficients(self, like: Array) -> dict[str, Array]: + """The coefficients of the kernel formula, derived from its parameters. + + Parameters + ---------- + like + An array giving the namespace, dtype and device of the coefficients. + + Returns + ------- + dict[str, Array] + Every coefficient with its own shape, differentiable with respect + to the parameters. + """ + raise NotImplementedError + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + """Evaluate the kernel formula with given coefficients. + + Parameters + ---------- + r_sq + Squared distances with shape ``(P,)``, in Ų. + coefficients + The coefficients of :meth:`coefficients`, each with its own shape + or with a leading pair axis of length ``P``. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + raise NotImplementedError + + def kspace(self, k_sq: Array) -> Array: + r"""Fourier transform of the kernel on squared wave numbers. + + :math:`\hat K(k) = \int K(r) e^{-i \mathbf{k} \cdot \mathbf{r}} d^3r`, + so that the energy of charges in a periodic cell of volume ``V`` reads + :math:`(k_e / 2V) \sum_{\mathbf{k} \neq 0} \hat K(k) |\hat\rho(\mathbf{k})|^2`. + + Parameters + ---------- + k_sq + Squared wave numbers with shape ``(nk,)`` or ``(rows, nk)``, in + Å⁻²; every entry must be positive. + + Returns + ------- + Array + The transformed kernel with the shape of ``k_sq``, in ų. + """ + raise NotImplementedError + + def contact_value(self, like: Array) -> Array: + r"""The kernel at contact, :math:`K(0)`, in Å⁻¹. + + The k-space lattice sum of a periodic system contains the self terms, + which the evaluation subtracts as :math:`\frac{1}{2} K(0) \sum_i q_i^2`. + + Parameters + ---------- + like + An array giving the namespace, dtype and device of the result. + """ + raise NotImplementedError + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + """Wave number at which the kernel weight falls below ``tolerance``. + + The Gaussian tail of the Fourier-transformed kernel makes the k-space + truncation exponentially convergent; the k grid of a periodic + evaluation ends at this wave number. + + Parameters + ---------- + tolerance + The relative weight of the slowest-decaying mode at the cutoff. + """ + raise NotImplementedError + + def _parameter(self, name: str, like: Array) -> Array: + """A parameter in the namespace, dtype and device of ``like``, autograd kept.""" + xp = array_api_compat.array_namespace(like) + return xp_asarray_nodetach( + xp, + getattr(self, name), + dtype=like.dtype, + device=array_api_compat.device(like), + ) + + @staticmethod + def _row_values(value: Array, k_sq: Array) -> Array: + """A scalar broadcast to shape ``(..., 1)`` over the leading axes of ``k_sq``. + + The gradient of a parameter of the transform sums over every wave + number. Broadcast over the rows, it sums along every row first, a + reduction a compiled graph spreads over the rows, and then over the + rows, rather than over every wave number in one block. + """ + xp = array_api_compat.array_namespace(k_sq) + return xp.broadcast_to(value, (*k_sq.shape[:-1], 1)) + + def serialize(self) -> dict: + return { + "@class": "InteractionKernel", + "type": self.TYPE, + "@version": 1, + "trainable": self.trainable, + "precision": self.precision, + "@variables": { + name: to_numpy_array(getattr(self, name)) + for name in self.PARAMETER_NAMES + }, + } + + @classmethod + def deserialize(cls, data: dict) -> "BaseInteractionKernel": + if cls is BaseInteractionKernel: + return cls.get_class_by_type(data["type"]).deserialize(data) + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + data.pop("type") + variables = data.pop("@variables") + obj = cls(**data, **cls._shape_kwargs(variables)) + for name in cls.PARAMETER_NAMES: + setattr(obj, name, np.asarray(variables[name], dtype=obj.prec)) + return obj + + @classmethod + def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: + """Constructor arguments that size the parameters stored in ``variables``.""" + return {} + + +# Taylor coefficients of erf(x) / x in powers of x^2, 2 / sqrt(pi) (-1)^n / +# (n! (2n + 1)), and the square of x below which LESKernel evaluates them: there +# the nine terms reach the float64 rounding of erf(x) / x. Above it the closed +# form keeps the relative error of its radial derivative within 2e-14 in +# float64 and 1e-5 in float32. +_ERF_OVER_X_SERIES = ( + 1.1283791670955126, + -0.37612638903183754, + 0.11283791670955126, + -0.026866170645131252, + 0.005223977625442188, + -0.0008548327023450852, + 0.00012055332981789664, + -1.492565035840625e-05, + 1.6462114365889246e-06, +) +_ERF_SERIES_X_SQ = 0.0625 + + +@BaseInteractionKernel.register("les") +class LESKernel(BaseInteractionKernel): + r"""Latent Ewald summation kernel :math:`K(r) = \mathrm{erf}(\alpha r) / r`. + + The kernel is the interaction of two Gaussian charge clouds of width + :math:`1/(\sqrt{2}\alpha)`: it reaches :math:`1/r` beyond a few + :math:`1/\alpha` and stays finite, :math:`2\alpha/\sqrt{\pi}`, at contact. + + Parameters + ---------- + alpha + The inverse width :math:`\alpha > 0`, in Å⁻¹. + trainable + Whether :math:`\alpha` is trained. + precision + The precision in which :math:`\log \alpha` is stored. + """ + + TYPE = "les" + PARAMETER_NAMES = ("log_alpha",) + + def __init__( + self, alpha: float = 1.0, trainable: bool = True, precision: str = "default" + ) -> None: + super().__init__(trainable, precision) + if not np.isfinite(alpha) or alpha <= 0.0: + raise ValueError(f"`alpha` must be finite and positive, got {alpha}") + self.log_alpha = np.array(np.log(alpha), dtype=self.prec) + + def coefficients(self, like: Array) -> dict[str, Array]: + xp = array_api_compat.array_namespace(like) + return {"alpha": xp.exp(self._parameter("log_alpha", like))} + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + xp = array_api_compat.array_namespace(r_sq) + alpha = coefficients["alpha"] + x_sq = alpha**2 * r_sq + # K = alpha erf(x) / x with x = alpha r. The derivative of the closed + # form subtracts two terms that agree to second order in x, so near + # contact erf(x) / x comes from its Taylor series in x^2. Both are + # evaluated on every pair, which keeps the shapes static; each reads a + # value of its own domain where the other is selected, which keeps the + # derivatives of both finite. + near = x_sq < _ERF_SERIES_X_SQ + x_far = xp.sqrt(xp.where(near, xp.ones_like(x_sq), x_sq)) + x_sq_near = xp.where(near, x_sq, xp.zeros_like(x_sq)) + series = xp.full_like(x_sq, _ERF_OVER_X_SERIES[-1]) + for coefficient in reversed(_ERF_OVER_X_SERIES[:-1]): + series = series * x_sq_near + coefficient + return alpha * xp.where(near, series, xp_erf(x_far) / x_far) + + def kspace(self, k_sq: Array) -> Array: + xp = array_api_compat.array_namespace(k_sq) + # erf(alpha r)/r is the interaction of two Gaussian charge clouds of + # width 1/(sqrt(2) alpha); its transform is the damped Coulomb + # 4 pi exp(-k^2 / 4 alpha^2) / k^2. + alpha = self._row_values(xp.exp(self._parameter("log_alpha", k_sq)), k_sq) + return 4.0 * np.pi * xp.exp(-k_sq / (4.0 * alpha * alpha)) / k_sq + + def contact_value(self, like: Array) -> Array: + xp = array_api_compat.array_namespace(like) + alpha = xp.exp(self._parameter("log_alpha", like)) + return 2.0 * alpha / np.sqrt(np.pi) + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + # exp(-k^2 / 4 alpha^2) < tolerance + alpha = float(np.exp(to_numpy_array(self.log_alpha).astype(np.float64))) + return 2.0 * alpha * np.sqrt(np.log(1.0 / tolerance)) + + +@BaseInteractionKernel.register("sog") +class SOGKernel(BaseInteractionKernel): + r"""Sum-of-Gaussians kernel :math:`K(r) = \sum_k a_k \exp(-r^2 / 2\sigma_k^2)`. + + Parameters + ---------- + amplitude + The amplitudes :math:`a_k`, in Å⁻¹, of either sign: a scalar shared by + every Gaussian or one value per Gaussian. Required with ``width``; + with the default widths it defaults to + :math:`a_k = \sqrt{2/\pi} \ln b / \sigma_k`, the trapezoidal + discretization in :math:`\ln s` of + :math:`1/r = \sqrt{2/\pi} \int_0^\infty e^{-r^2/2s^2} s^{-2} \, ds`. + The default kernel then follows :math:`1/r`, less a nearly constant + offset, from a few :math:`\sigma` to a fraction of + :math:`\sigma b^{M-1}`, and stays finite at contact. The offset + exerts no force between fixed charges; through charges that respond + to the positions it does. + width + The widths :math:`\sigma_k > 0`, in Å. Defaults to the geometric + sequence :math:`\sigma_k = \sigma b^k`, :math:`k = 0, \dots, M - 1`. + n_gaussian + The number :math:`M` of Gaussians of the default widths. + base_width + The width :math:`\sigma` of the first default Gaussian, in Å. + width_ratio + The ratio :math:`b > 1` of consecutive default widths. + trainable + Whether the amplitudes and widths are trained. + precision + The precision in which the amplitudes and log-widths are stored. + """ + + TYPE = "sog" + PARAMETER_NAMES = ("amplitude", "log_width") + + def __init__( + self, + amplitude: float | list[float] | None = None, + width: list[float] | None = None, + n_gaussian: int = 12, + base_width: float = 2.180230445405648, + width_ratio: float = 1.62976708826776469, + trainable: bool = True, + precision: str = "default", + ) -> None: + super().__init__(trainable, precision) + if width is None: + if n_gaussian < 1: + raise ValueError(f"`n_gaussian` must be positive, got {n_gaussian}") + if not np.isfinite(base_width) or base_width <= 0.0: + raise ValueError(f"`base_width` must be positive, got {base_width}") + if not np.isfinite(width_ratio) or width_ratio <= 1.0: + raise ValueError( + f"`width_ratio` must be greater than one, got {width_ratio}" + ) + width = base_width * np.power( + width_ratio, np.arange(n_gaussian, dtype=np.float64) + ) + if amplitude is None: + amplitude = np.sqrt(2.0 / np.pi) * np.log(width_ratio) / width + elif amplitude is None: + raise ValueError("`amplitude` is required with explicit `width`") + width = np.asarray(width, dtype=np.float64).reshape(-1) + if width.size < 1 or not np.all(np.isfinite(width)) or np.any(width <= 0.0): + raise ValueError("`width` must hold finite positive values") + amplitude = np.asarray(amplitude, dtype=np.float64).reshape(-1) + if amplitude.size == 1: + amplitude = np.full_like(width, amplitude[0]) + if amplitude.shape != width.shape or not np.all(np.isfinite(amplitude)): + raise ValueError( + "`amplitude` must be a finite scalar or hold one value per Gaussian" + ) + self.amplitude = amplitude.astype(self.prec) + self.log_width = np.log(width).astype(self.prec) + + def coefficients(self, like: Array) -> dict[str, Array]: + xp = array_api_compat.array_namespace(like) + return { + "amplitude": self._parameter("amplitude", like), + "inv_two_width_sq": 0.5 * xp.exp(-2.0 * self._parameter("log_width", like)), + } + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + xp = array_api_compat.array_namespace(r_sq) + gaussian = xp.exp(-r_sq[:, None] * coefficients["inv_two_width_sq"]) + return xp.sum(coefficients["amplitude"] * gaussian, axis=-1) + + def kspace(self, k_sq: Array) -> Array: + xp = array_api_compat.array_namespace(k_sq) + amplitude = self._parameter("amplitude", k_sq) + width = xp.exp(self._parameter("log_width", k_sq)) + # FT of a_k exp(-r^2 / 2 sigma_k^2) is a_k (2 pi)^{3/2} sigma_k^3 + # exp(-sigma_k^2 k^2 / 2); for the default amplitudes + # a_k = sqrt(2/pi) ln b / sigma_k the coefficient is 4 pi ln b sigma_k^2, + # whose ln-sigma trapezoid sum approaches 4 pi / k^2. The mode loop is + # unrolled: a (nk, M) intermediate breaks inductor's double backward of + # the k-space virial (NaN), while per-mode (nk,) terms compile fine. + terms = [] + for m in range(width.shape[0]): + weight = self._row_values(amplitude[m] * width[m] ** 3, k_sq) + rate = self._row_values(0.5 * width[m] ** 2, k_sq) + terms.append((2.0 * np.pi) ** 1.5 * weight * xp.exp(-rate * k_sq)) + return sum(terms[1:], terms[0]) + + def contact_value(self, like: Array) -> Array: + xp = array_api_compat.array_namespace(like) + amplitude = self._parameter("amplitude", like) + return xp.sum(amplitude) + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + # the slowest-decaying mode is the NARROWEST Gaussian: + # exp(-sigma_min^2 k^2 / 2) < tolerance + sigma_min = float( + np.exp(np.min(to_numpy_array(self.log_width).astype(np.float64))) + ) + return np.sqrt(2.0 * np.log(1.0 / tolerance)) / sigma_min + + @classmethod + def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: + return {"n_gaussian": int(np.asarray(variables["log_width"]).size)} diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py new file mode 100644 index 0000000000..3795fb8167 --- /dev/null +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -0,0 +1,730 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import math +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.common import ( + GLOBAL_NP_FLOAT_PRECISION, + NativeOP, + to_numpy_array, +) +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + K_GRID, + PME_MESH, + frame_id_from_n_node, + k_grid_bounds_of_cells, + segment_sum, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, +) +from .reciprocal import ( + kspace_energy, + mesh_energy, +) + +CHARGE_TARGET_SOURCES = ("neutral", "charge_spin", "fparam") +"""Sources of the total charge a constrained channel sums to.""" + +CHARGE_NET_NEURON = [64, 64] +"""Default hidden widths of the charge network.""" + +PME_ORDER = 8 +"""Default order of the B-splines that spread the charges to the mesh.""" + +PME_SPACING_FRACTION = 0.75 +"""Default mesh spacing, as a fraction of the Nyquist spacing of the wave-number cutoff. + +At order 8 the mesh evaluation of a periodic frame then stays within a few +parts in a million of the sum over the wave numbers, in the reciprocal energy +and in the force. +""" + +CHARGE_NET_SEED_INDEX = 10_000 +"""Child index of the fitting seed that seeds a charge network configured without a seed. + +The networks of the fitting and of the descriptor draw from small child +indices of their seeds, so the charge network draws a stream of its own even +when the fitting and descriptor seeds coincide. +""" + + +@BaseLongRange.register("latent_charge") +class LatentChargeLongRange(NativeOP, BaseLongRange): + r"""Energy of interacting latent charges, periodic or not. + + A charge network predicts :math:`C` latent charge channels + :math:`q_i^c` per atom, in units of :math:`e`, from the descriptor shared + with the energy fitting; padding and atoms owned by another rank carry no + charge. For a non-periodic batch the energy is the direct sum over the + pairs of atoms of each frame listed by the + :data:`~deepmd.dpmodel.utils.neighbor_graph.ALL_PAIRS` topology, each + unordered pair once; for a periodic batch (with ``periodic=True``) it is + the lattice sum of the kernel over the images of the cell, evaluated in + reciprocal space (see :mod:`~deepmd.dpmodel.long_range.reciprocal`), + + .. math:: + E_{\rm non-periodic} = k_e \sum_{i None: + if charge_net.vacuum_ref: + raise ValueError( + "the charge network reads the descriptor of the real atoms only " + "and cannot reference isolated atoms (`vacuum_ref`)" + ) + self.charge_net = charge_net + self.kernel = kernel + self.charge_constraint = self._normalize_constraint( + charge_constraint, charge_net + ) + self.periodic = bool(periodic) + self.kspace_tolerance = float(kspace_tolerance) + if ( + not np.isfinite(self.kspace_tolerance) + or not 0.0 < self.kspace_tolerance < 1.0 + ): + raise ValueError( + f"`kspace_tolerance` must be in (0, 1), got {self.kspace_tolerance}" + ) + # The frame charge/spin condition a compressed model is specialized + # to; see :meth:`fold_charge_state`. + self.charge_state = None + if k_grid_mode is None: + k_grid_mode = "pme" if self.periodic else "ragged" + self.k_grid_mode = str(k_grid_mode) + if self.k_grid_mode not in ("ragged", "fixed", "pme"): + raise ValueError( + f"`k_grid_mode` must be 'ragged', 'fixed' or 'pme', got {self.k_grid_mode!r}" + ) + if self.k_grid_mode in ("fixed", "pme") and not self.periodic: + raise ValueError( + f"`k_grid_mode='{self.k_grid_mode}'` requires `periodic=True`" + ) + if k_grid_bounds is not None and self.k_grid_mode != "fixed": + raise ValueError( + "`k_grid_bounds` is only meaningful with `k_grid_mode='fixed'`" + ) + self.decouple_k_grid = bool(decouple_k_grid) + if self.decouple_k_grid and not self.periodic: + raise ValueError("`decouple_k_grid` requires `periodic=True`") + # A decoupled grid reads the kernel once, at construction: the k-point + # selection stays at the initial kernel however training moves it, + # while the kernel values in the graph still use the live parameters. + self.frozen_k_max = ( + float(self.kernel.kspace_cutoff(self.kspace_tolerance)) + if self.decouple_k_grid + else None + ) + # The per-axis bounds of the shared integer grid of the fixed mode, + # zero until the data statistics size it (see + # :meth:`compute_input_stats`): an array of fixed shape, so that a + # backend keeping arrays in its state dict sees the same keys and + # shapes before and after the sizing, and a checkpoint restores the + # sized grid. + if k_grid_bounds is not None: + bounds = np.asarray(k_grid_bounds, dtype=np.int64).reshape(-1) + if bounds.shape != (3,) or np.any(bounds < 0) or not np.any(bounds): + raise ValueError( + "`k_grid_bounds` must be three non-negative integers, not " + f"all zero, got {list(k_grid_bounds)}" + ) + self.k_grid_bounds = bounds + elif self.k_grid_mode == "fixed": + self.k_grid_bounds = np.zeros(3, dtype=np.int64) + else: + self.k_grid_bounds = None + # The mesh of a periodic frame, on which training in the "pme" mode + # and every frozen model evaluate the reciprocal sum. + if not self.periodic and (pme_spacing is not None or pme_order is not None): + raise ValueError("`pme_spacing` and `pme_order` require `periodic=True`") + self.pme_order = PME_ORDER if pme_order is None else int(pme_order) + if self.pme_order < 2: + raise ValueError(f"`pme_order` must be at least 2, got {self.pme_order}") + if pme_spacing is not None: + pme_spacing = float(pme_spacing) + if not np.isfinite(pme_spacing) or pme_spacing <= 0.0: + raise ValueError( + f"`pme_spacing` must be finite and positive, got {pme_spacing}" + ) + self.pme_spacing = pme_spacing + + @staticmethod + def _normalize_constraint( + constraint: dict | None, charge_net: LatentChargeFitting + ) -> dict | None: + if constraint is None: + return None + source = str(constraint.get("source", "neutral")) + channel = int(constraint.get("channel", 0)) + fparam_index = int(constraint.get("fparam_index", 0)) + if source not in CHARGE_TARGET_SOURCES: + raise ValueError( + f"charge constraint source must be one of {CHARGE_TARGET_SOURCES}, " + f"got {source!r}" + ) + if not 0 <= channel < charge_net.dim_out: + raise ValueError( + f"constrained channel {channel} is out of range for " + f"{charge_net.dim_out} latent charge channel(s)" + ) + if source == "fparam" and not 0 <= fparam_index < charge_net.numb_fparam: + raise ValueError( + f"`fparam_index` {fparam_index} is out of range for " + f"{charge_net.numb_fparam} frame parameter(s)" + ) + return {"source": source, "channel": channel, "fparam_index": fparam_index} + + @classmethod + def from_config( + cls, + config: dict, + descriptor: Any, + fitting: Any, + fitting_seed: int | list[int] | None = None, + ) -> "LatentChargeLongRange": + """Build the module from its configuration. + + The charge network reads the descriptor and takes the conditioning + inputs of the energy fitting: its frame and atomic parameters, their + defaults and its case embedding. It follows the activation function, + skip connection and precision of the fitting unless its configuration + sets them, and its seed defaults to a child of the fitting seed. The + kernel parameters are stored in the precision of the charge network. + + Parameters + ---------- + config + The module configuration: ``channels``, ``charge_net``, ``kernel`` + (the ``les`` kernel when absent) and ``charge_constraint``. + descriptor + The descriptor of the model. + fitting + The energy fitting of the model. + fitting_seed + The seed of the energy fitting. + """ + charge_net_config = { + "neuron": CHARGE_NET_NEURON, + "activation_function": fitting.activation_function, + "resnet_dt": fitting.resnet_dt, + "precision": fitting.precision, + "seed": child_seed(fitting_seed, CHARGE_NET_SEED_INDEX), + **config.get("charge_net", {}), + } + kernel_config = { + "type": "les", + "precision": charge_net_config["precision"], + **config.get("kernel", {}), + } + kernel_cls = BaseInteractionKernel.get_class_by_type(kernel_config.pop("type")) + charge_net = LatentChargeFitting( + var_name="latent_charge", + ntypes=descriptor.get_ntypes(), + dim_descrpt=descriptor.get_dim_out(), + dim_out=int(config.get("channels", 1)), + numb_fparam=fitting.get_dim_fparam(), + numb_aparam=fitting.get_dim_aparam(), + dim_case_embd=fitting.dim_case_embd, + default_fparam=fitting.get_default_fparam(), + use_aparam_as_mask=fitting.use_aparam_as_mask, + mixed_types=descriptor.mixed_types(), + type_map=fitting.get_type_map(), + **charge_net_config, + ) + return cls( + charge_net, + kernel_cls(**kernel_config), + config.get("charge_constraint"), + periodic=config.get("periodic", False), + kspace_tolerance=config.get("kspace_tolerance", 1e-8), + k_grid_mode=config.get("k_grid_mode"), + k_grid_bounds=config.get("k_grid_bounds"), + decouple_k_grid=config.get("decouple_k_grid", False), + pme_spacing=config.get("pme_spacing"), + pme_order=config.get("pme_order"), + ) + + def _k_max(self) -> float: + """The wave-number cutoff selecting the k points of a periodic batch. + + A coupled grid reads the current kernel parameters, so a trained + kernel re-derives its grid every batch; a decoupled grid keeps the + cutoff of the initial parameters. + """ + if self.decouple_k_grid: + return self.frozen_k_max + return float(self.kernel.kspace_cutoff(self.kspace_tolerance)) + + def mesh_spacing(self) -> float: + r"""The largest spacing of the mesh of a periodic frame, in Å. + + The configured ``pme_spacing``, or :data:`PME_SPACING_FRACTION` of the + Nyquist spacing :math:`\pi / k_{max}` of the wave-number cutoff. + """ + if self.pme_spacing is not None: + return self.pme_spacing + return PME_SPACING_FRACTION * math.pi / self._k_max() + + def use_mesh_evaluation(self) -> None: + """Evaluate the reciprocal sum of every periodic batch on the mesh. + + A frozen model takes the mesh of every frame as an input of dynamic + size, sized by the cell of the frame, so that it serves any cell at + the accuracy of :meth:`mesh_spacing`; the explicit wave numbers of + the other modes have a count that no input layout fixes. The spacing + is fixed at the current kernel parameters, a constant of the frozen + model. + """ + if self.periodic: + self.pme_spacing = self.mesh_spacing() + self.k_grid_mode = "pme" + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the reciprocal-space topology of a periodic batch. + + The wave-number cutoff reads the kernel parameters at build time, so + that a trained kernel moves the k grid and the mesh with it, unless + ``decouple_k_grid`` keeps the initial cutoff. The k grid holds the + wave numbers of every frame within the cutoff, or the points of the + fixed integer grid with those beyond it masked; the mesh resolves + every cell of the batch at :meth:`mesh_spacing`. + + Raises + ------ + ValueError + If the fixed grid has not been sized. + """ + if not self.periodic: + return {} + if self.k_grid_mode == "pme": + return {PME_MESH: {"spacing": self.mesh_spacing()}} + params = {"k_max": self._k_max()} + if self.k_grid_mode == "fixed": + if not self._k_grid_sized(): + raise ValueError( + "the fixed k grid is not sized yet: run the data " + "statistics (`compute_input_stats`) on periodic " + "frames, or configure `k_grid_bounds`" + ) + params["bounds"] = to_numpy_array(self.k_grid_bounds) + return {K_GRID: params} + + def required_topologies(self) -> tuple[str, ...]: + if self.periodic: + if self.k_grid_mode == "pme": + return (ALL_PAIRS, PME_MESH) + return (ALL_PAIRS, K_GRID) + return (ALL_PAIRS,) + + def uses_charge_spin(self) -> bool: + # A folded state stands in for the condition. + return ( + self.charge_constraint is not None + and self.charge_constraint["source"] == "charge_spin" + and self.charge_state is None + ) + + def fold_charge_state(self, charge_spin: list[float]) -> None: + """Specialize the constraint to one frame charge/spin condition. + + A compressed descriptor folds the condition into its tables and reads + none at run time, so a constraint that reads its total charge from the + condition takes it from this state instead. The state is a constant + of the deployed model that a deployment rewrites, together with the + tables of the descriptor, when it serves another condition. + + Parameters + ---------- + charge_spin + The condition ``[charge, spin multiplicity]``. + """ + self.charge_state = np.asarray( + charge_spin, dtype=GLOBAL_NP_FLOAT_PRECISION + ).reshape(1, 2) + + def charge_state_constants(self) -> list[Array]: + return [] if self.charge_state is None else [self.charge_state] + + def rebuild_charge_state(self, charge_spin: Array) -> list[Array]: + if self.charge_state is None: + return [] + xp = array_api_compat.array_namespace(charge_spin) + state = xp.astype(charge_spin, self.charge_state.dtype) + return [xp.reshape(state, (1, 2))] + + def call(self, *args: Any, **kwargs: Any) -> Array: + return self.call_graph(*args, **kwargs) + + def call_graph( + self, + descriptor: Array, + atype: Array, + node_mask: Array, + n_node: Array, + topologies: dict[str, Any], + fparam: Array | None = None, + aparam: Array | None = None, + charge_spin: Array | None = None, + ) -> Array: + xp = array_api_compat.array_namespace(descriptor) + n = descriptor.shape[0] + pairs = topologies[ALL_PAIRS] + frame_id = frame_id_from_n_node(n_node, n_total=n) + fparam_node = None if fparam is None else xp.take(fparam, frame_id, axis=0) + charge = self.charge_net.call_graph( + descriptor, atype, fparam=fparam_node, aparam=aparam + )["latent_charge"] + # The pair sum cancels between charges of opposite sign, so it runs in + # the precision of the geometry. + charge = xp.astype(charge, pairs.vec.dtype) + charge = xp.where(node_mask[:, None], charge, xp.zeros_like(charge)) + if self.charge_constraint is not None: + charge = self._constrain( + charge, node_mask, frame_id, n_node, fparam, charge_spin + ) + energy = self._pair_branch(charge, pairs, n) + if self.periodic: + # The lattice sum of a periodic frame and its self term cancel to + # a small fraction of either, so the reciprocal-space energy is + # formed in float64 whatever the precision of the geometry. + charge_64 = xp.astype(charge, xp.float64) + if self.k_grid_mode == "pme": + reciprocal = mesh_energy( + charge_64, + topologies[PME_MESH], + frame_id, + self.kernel, + self.pme_order, + ) + else: + reciprocal = kspace_energy(charge_64, topologies[K_GRID], self.kernel) + energy = energy + xp.astype(reciprocal, energy.dtype) + return xp.reshape(energy, (n, 1)) + + def _pair_branch(self, charge: Array, pairs: Any, n: int) -> Array: + """Direct pair sum over the non-periodic frames, per atom.""" + xp = array_api_compat.array_namespace(charge) + src = pairs.index[0] + dst = pairs.index[1] + pair_charge = xp.sum( + xp.take(charge, src, axis=0) * xp.take(charge, dst, axis=0), axis=-1 + ) + r_sq = xp.sum(pairs.vec * pairs.vec, axis=-1) + pair_energy = pair_charge * self._pair_kernel(r_sq, src, n) + half_energy = (0.5 * COULOMB_CONSTANT) * xp.where( + pairs.mask, pair_energy, xp.zeros_like(r_sq) + ) + return segment_sum(half_energy, src, n) + segment_sum(half_energy, dst, n) + + def _pair_kernel(self, r_sq: Array, src: Array, n: int) -> Array: + """Kernel values of the pairs. + + The pairs are located on the node axis for a backend that brings the + kernel coefficients to the pairs through it. + + Parameters + ---------- + r_sq + Squared pair distances with shape ``(P,)``, in Ų. + src + First node of every pair with shape ``(P,)``. + n + The number of nodes. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + return self.kernel(r_sq) + + def _constrain( + self, + charge: Array, + node_mask: Array, + frame_id: Array, + n_node: Array, + fparam: Array | None, + charge_spin: Array | None, + ) -> Array: + """Shift the constrained channel so that it sums to the frame total charge.""" + xp = array_api_compat.array_namespace(charge) + nf = n_node.shape[0] + channel = self.charge_constraint["channel"] + target = self._total_charge(nf, charge, fparam, charge_spin) + weight = xp.astype(node_mask, charge.dtype) + count = segment_sum(weight, frame_id, nf) + total = segment_sum(charge[:, channel], frame_id, nf) + # A frame without a charged atom has nothing to shift. + shift = (total - target) / xp.maximum(count, xp.ones_like(count)) + constrained = charge[:, channel] - xp.take(shift, frame_id, axis=0) * weight + return xp.concat( + [charge[:, :channel], constrained[:, None], charge[:, channel + 1 :]], + axis=-1, + ) + + def _total_charge( + self, + nf: int, + like: Array, + fparam: Array | None, + charge_spin: Array | None, + ) -> Array: + """The total charge of every frame, which the constrained channel sums to. + + Parameters + ---------- + nf + The number of frames. + like + An array giving the namespace, dtype and device of the result. + fparam, charge_spin + As in :meth:`call_graph`. + + Returns + ------- + Array + The total charges with shape ``(nf,)``, in units of :math:`e`. + """ + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + source = self.charge_constraint["source"] + if source == "charge_spin" and self.charge_state is not None: + state = xp.asarray(self.charge_state, device=device) + target = xp.broadcast_to(state[:, 0], (nf,)) + elif source == "charge_spin": + if charge_spin is None: + raise ValueError( + "the charge constraint reads the total charge from " + "`charge_spin`, which was not given" + ) + target = charge_spin[:, 0] + elif source == "fparam": + if fparam is None: + raise ValueError( + "the charge constraint reads the total charge from " + "`fparam`, which was not given" + ) + target = fparam[:, self.charge_constraint["fparam_index"]] + else: + target = xp.zeros((nf,), dtype=like.dtype, device=device) + return xp.astype(target, like.dtype) + + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: + if self.k_grid_mode == "fixed" and not self._k_grid_sized(): + # Materialize the samples once for the grid sizing and the + # networks; a lazy sampler is shared by the callers. + sampled = merged() if callable(merged) else merged + self._size_k_grid(sampled) + merged = sampled + self.charge_net.compute_input_stats( + merged, protection=protection, stat_file_path=stat_file_path + ) + + def _size_k_grid(self, sampled: list[dict]) -> None: + r"""Size the fixed integer k grid by the largest sampled cell. + + Every sampled frame needs :math:`|n_d| \le k_{max} |a_d| / 2\pi` per + axis at the wave-number cutoff of the kernel; the grid takes the + elementwise maximum of the bounds over the frames. The kernel + parameters are read at statistics time, so the grid is sized for the + kernel of that moment; frames that outgrow it later, by larger cells + or a kernel whose cutoff grows, are truncated to the grid. + """ + k_max = self._k_max() + bounds = None + for frame in sampled: + box = frame.get("box") + if box is None: + continue + box = np.asarray(to_numpy_array(box), dtype=np.float64) + if box.size == 0: + continue + frame_max = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + bounds = frame_max if bounds is None else np.maximum(bounds, frame_max) + if bounds is None: + raise ValueError( + "the fixed k grid takes its size from the data statistics, " + "but the sampled frames hold no simulation cell; provide " + "periodic training data or configure `k_grid_bounds`" + ) + self.k_grid_bounds = bounds.astype(np.int64) + + def _k_grid_sized(self) -> bool: + """Whether the fixed integer grid has its bounds.""" + return bool(np.any(to_numpy_array(self.k_grid_bounds))) + + def change_type_map(self, type_map: list[str]) -> None: + self.charge_net.change_type_map(type_map=type_map) + + def set_case_embd(self, case_idx: int) -> None: + self.charge_net.set_case_embd(case_idx) + + def serialize(self) -> dict: + data = { + "@class": "LongRange", + "type": "latent_charge", + "@version": 1, + "charge_net": self.charge_net.serialize(), + "kernel": self.kernel.serialize(), + "charge_constraint": self.charge_constraint, + } + if self.periodic: + data["periodic"] = True + data["kspace_tolerance"] = self.kspace_tolerance + data["k_grid_mode"] = self.k_grid_mode + if self.k_grid_mode == "fixed" and self._k_grid_sized(): + data["k_grid_bounds"] = to_numpy_array(self.k_grid_bounds).tolist() + if self.decouple_k_grid: + data["decouple_k_grid"] = True + data["frozen_k_max"] = self.frozen_k_max + if self.periodic: + data["pme_order"] = self.pme_order + if self.pme_spacing is not None: + data["pme_spacing"] = self.pme_spacing + if self.charge_state is not None: + data["charge_state"] = to_numpy_array(self.charge_state).tolist() + return data + + @classmethod + def deserialize(cls, data: dict) -> "LatentChargeLongRange": + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + obj = cls( + LatentChargeFitting.deserialize(data["charge_net"]), + BaseInteractionKernel.deserialize(data["kernel"]), + data["charge_constraint"], + periodic=data.get("periodic", False), + kspace_tolerance=data.get("kspace_tolerance", 1e-8), + # A periodic record without a mode is one of the explicit wave + # numbers. + k_grid_mode=data.get("k_grid_mode", "ragged"), + k_grid_bounds=data.get("k_grid_bounds"), + decouple_k_grid=data.get("decouple_k_grid", False), + pme_spacing=data.get("pme_spacing"), + pme_order=data.get("pme_order"), + ) + if data.get("frozen_k_max") is not None: + # The constructor derives the frozen cutoff from the (possibly + # trained) kernel; the record holds the value the grid was + # actually built with. + obj.frozen_k_max = float(data["frozen_k_max"]) + if data.get("charge_state") is not None: + obj.fold_charge_state(np.asarray(data["charge_state"]).reshape(-1).tolist()) + return obj diff --git a/deepmd/dpmodel/long_range/reciprocal.py b/deepmd/dpmodel/long_range/reciprocal.py new file mode 100644 index 0000000000..27b74e00e4 --- /dev/null +++ b/deepmd/dpmodel/long_range/reciprocal.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Reciprocal-space energy of latent charges in periodic cells. + +For the charges :math:`q_i^c` of a periodic cell of volume :math:`V` and a +pair kernel :math:`K` with Fourier transform :math:`\hat K`, the energy of the +lattice of images, + +.. math:: + E = k_e \Big[ \frac{1}{2V} \sum_{\mathbf{k} \neq \mathbf{0}} \sum_c + \hat K(\mathbf{k}) \, |\hat\rho_c(\mathbf{k})|^2 + - \frac{1}{2} K(0) \sum_{i,c} (q_i^c)^2 \Big], + \qquad + \hat\rho_c(\mathbf{k}) = \sum_i q_i^c e^{i \mathbf{k} \cdot \mathbf{r}_i}, + +is the sum over the wave numbers of the cell minus the interaction of every +charge with itself. The :math:`\mathbf{k} = \mathbf{0}` mode is dropped, the +tin-foil convention, which for a channel with a net charge amounts to a +uniform neutralizing background. The energy is attributed to the atoms as +:math:`E_i = \frac{1}{2} \sum_c q_i^c \phi_i^c`, with the potential +:math:`\phi_i^c = \partial E / \partial q_i^c` that the other charges and the +images create at atom :math:`i`, as the pair sum of a non-periodic system +splits every pair term evenly between its atoms. + +Two evaluations serve the topologies of +:mod:`~deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph`: +:func:`kspace_energy` sums over explicit wave numbers, and +:func:`mesh_energy` evaluates the same sum by the smooth particle-mesh Ewald +method on a mesh. +""" + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, + xp_add_at, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, + MeshGrid, +) + +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, +) + + +def _self_potential( + kernel: BaseInteractionKernel, charge: Array, periodic_node: Array +) -> Array: + """The potential of every charge at itself, k_e K(0) q, on periodic nodes.""" + xp = array_api_compat.array_namespace(charge) + return ( + COULOMB_CONSTANT + * kernel.contact_value(charge) + * charge + * xp.astype(periodic_node, charge.dtype)[:, None] + ) + + +def _strained_rows(xp: Any, rows: Array, strain: Array) -> Array: + r"""Wave numbers of the strained cells, to first order in the strain. + + The strain :math:`\varepsilon` of a frame maps its cell :math:`B` to + :math:`B S^\mathsf{T}`, :math:`S = I + \varepsilon`, and a wave number, + a row :math:`\mathbf{k}`, to :math:`\mathbf{k} S^{-1}`. The strain of a + topology is zero, and the energy is differentiated once with respect to + it, for the virial; the value and the first derivative at zero strain + are those of :math:`S^{-1} = I - \varepsilon`, which maps the rows here + and keeps linear algebra out of the energy and its derivatives. + + Parameters + ---------- + rows + Wave numbers or reciprocal lattice vectors with shape ``(nf, ..., 3)``. + strain + The strain of every frame with shape ``(nf, 3, 3)``. + + Returns + ------- + Array + The strained wave numbers, of the shape of ``rows``. + """ + eps = xp.reshape(strain, (strain.shape[0],) + (1,) * (rows.ndim - 2) + (3, 3)) + return rows - xp.sum(rows[..., :, None] * eps, axis=-2) + + +def _strained_volume(volume: Array, strain: Array) -> Array: + r"""Volumes of the strained cells, :math:`V \det S = V (1 + \mathrm{tr}\,\varepsilon)` to first order.""" + return volume * (1.0 + strain[:, 0, 0] + strain[:, 1, 1] + strain[:, 2, 2]) + + +def kspace_energy( + charge: Array, + k_grid: KGrid, + kernel: BaseInteractionKernel, +) -> Array: + r"""Energy of the periodic frames over explicit wave numbers, per atom. + + The structure factors of the wave numbers of every frame come from the + phases of its nodes, and the potential at a node from the same phases: + :math:`\phi_i = \frac{k_e}{V} \sum_{\mathbf{k}} \hat K(\mathbf{k}) + [\cos(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Re}\,\hat\rho + + \sin(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Im}\,\hat\rho] - k_e K(0) q_i`. + Every frame pairs the wave numbers of its row of :attr:`KGrid.vec` with + the nodes of its row of :attr:`KGrid.slot`, in products batched over the + frames and the blocks of their wave numbers. The geometry is strained by + :attr:`KGrid.strain` (see :class:`KGrid`). + + Parameters + ---------- + charge + Node charges with shape ``(N, C)``, zero on the nodes that carry none. + k_grid + The wave numbers and the nodes of the periodic frames. + kernel + The pair kernel. + + Returns + ------- + Array + The energy of every node with shape ``(N,)``, in eV; zero in a frame + without wave numbers. + """ + xp = array_api_compat.array_namespace(charge) + device = array_api_compat.device(charge) + dtype = charge.dtype + n_node, n_channel = charge.shape + nf, n_slot = k_grid.slot.shape + # === Step 1. The nodes of every frame, (nf, ns, 3) and (nf, ns, C) === + real = k_grid.slot >= 0 + index = xp.reshape(xp.where(real, k_grid.slot, xp.zeros_like(k_grid.slot)), (-1,)) + coord = xp.take(xp.astype(k_grid.coord, dtype), index, axis=0) + coord = xp.reshape(coord, (nf, n_slot, 3)) + charge_slot = xp.reshape(xp.take(charge, index, axis=0), (nf, n_slot, n_channel)) + charge_slot = charge_slot * xp.astype(real, dtype)[:, :, None] + # === Step 2. Strained wave numbers and volumes: k' = S^-T k, V' = V det S === + # Every product over the wave numbers of a frame broadcasts the per-frame + # factor over the blocks, (nf, nb, kb, .), so that its sum runs over the + # points of a block. + strain = xp.astype(k_grid.strain, dtype) + k_vec = xp.astype(k_grid.vec, dtype) + strained_k = _strained_rows(xp, k_vec, strain) + volume = _strained_volume(xp.astype(k_grid.volume, dtype), strain) + # === Step 3. Structure factors, (nf, nb, kb, C) === + # The phase k' . r' = k . r does not change under the strain. + phase = xp.matmul(k_vec, xp.matrix_transpose(coord)[:, None, :, :]) + cos_phase = xp.cos(phase) + sin_phase = xp.sin(phase) + rho_real = xp.matmul(cos_phase, charge_slot[:, None, :, :]) + rho_imag = xp.matmul(sin_phase, charge_slot[:, None, :, :]) + # === Step 4. Potential at the nodes of every frame, (nf, ns, C) === + weight = COULOMB_CONSTANT * kernel.kspace(xp.sum(strained_k * strained_k, axis=-1)) + weight = weight / volume[:, None, None] + weight = xp.where(k_grid.mask, weight, xp.zeros_like(weight))[..., None] + block_potential = xp.matmul(xp.matrix_transpose(cos_phase), weight * rho_real) + block_potential = block_potential + xp.matmul( + xp.matrix_transpose(sin_phase), weight * rho_imag + ) + potential = xp.sum(block_potential, axis=1) + # The self term belongs to the lattice sum of a frame with wave numbers; + # a non-periodic batch evaluated by a periodic module carries none. + has_points = xp.astype(xp.any(k_grid.mask, axis=(1, 2)), dtype)[:, None, None] + potential = ( + potential + - COULOMB_CONSTANT * kernel.contact_value(charge) * charge_slot * has_points + ) + slot_energy = 0.5 * xp.sum(charge_slot * potential, axis=-1) + # === Step 5. Energy on the node axis === + # A padding slot carries no charge and adds a zero to the first node. + energy = xp.zeros((n_node,), dtype=dtype, device=device) + return xp_add_at(energy, index, xp.reshape(slot_energy, (-1,))) + + +def bspline_integer_weights(order: int) -> np.ndarray: + r"""The cardinal B-spline of order ``order`` at :math:`1, 2, \dots, p - 1`. + + These are the weights of the assignment factor + :math:`D(\theta) = \sum_{l=0}^{p-2} M_p(l + 1) e^{i \theta l}` of the + smooth particle-mesh Ewald method. + """ + return np.asarray( + [ + sum( + (-1.0) ** j * math.comb(order, j) * max(l + 1.0 - j, 0.0) ** (order - 1) + for j in range(order + 1) + ) + / math.factorial(order - 1) + for l in range(order - 1) + ] + ) + + +def _bspline(xp: Any, x: Array, order: int) -> Array: + r"""The cardinal B-spline :math:`M_p(x)` of order :math:`p`, elementwise. + + :math:`M_p(x) = \frac{1}{(p-1)!} \sum_{j=0}^{p} (-1)^j \binom{p}{j} + \max(x - j, 0)^{p-1}`, a piecewise polynomial supported on + :math:`[0, p]` with :math:`p - 2` continuous derivatives. + """ + value = None + for j in range(order + 1): + term = ((-1.0) ** j * math.comb(order, j)) * xp.maximum( + x - j, xp.zeros_like(x) + ) ** (order - 1) + value = term if value is None else value + term + return value / math.factorial(order - 1) + + +def _inverse3(xp: Any, m: Array) -> Array: + """Inverse of a batch of 3x3 matrices, its columns the cross products of the rows over the determinant.""" + a, b, c = m[..., 0, :], m[..., 1, :], m[..., 2, :] + bc = xp.linalg.cross(b, c) + ca = xp.linalg.cross(c, a) + ab = xp.linalg.cross(a, b) + det = xp.sum(a * bc, axis=-1) + return xp.stack([bc, ca, ab], axis=-1) / det[..., None, None] + + +def _stencil( + xp: Any, frac: Array, dims: tuple[Any, Any, Any], order: int +) -> tuple[Array, Array]: + r"""The mesh points every node spreads to and their weights. + + Along axis :math:`d` the scaled coordinate :math:`u = n_d s_d` of the + fractional coordinate :math:`s_d \in [0, 1)` reaches the ``order`` points + :math:`\lfloor u \rfloor - p + 1, \dots, \lfloor u \rfloor` with the + weights :math:`M_p(u - m)`; the stencil of a node is the outer product of + its three axes. + + Returns + ------- + tuple[Array, Array] + The flat mesh indices and the weights, both with shape + ``(N, order**3)``. + """ + device = array_api_compat.device(frac) + n_node = frac.shape[0] + shift = xp.arange(order, dtype=frac.dtype, device=device)[None, :] + offset = xp.arange(order, dtype=xp.int64, device=device)[None, :] + index = None + weight = None + for d, n_d in enumerate(dims): + u = frac[:, d] * n_d + base = xp.floor(u) + # The distance of point base - p + 1 + j from u is t + p - 1 - j. + w = _bspline(xp, (u - base)[:, None] + (order - 1) - shift, order) + m = (xp.astype(base, xp.int64)[:, None] - (order - 1) + offset) % n_d + if index is None: + index, weight = m, w + else: + index = xp.reshape(index[:, :, None] * n_d + m[:, None, :], (n_node, -1)) + weight = xp.reshape(weight[:, :, None] * w[:, None, :], (n_node, -1)) + return index, weight + + +def _assignment_factor( + m_axes: list[Array], dims: tuple[Any, Any, Any], order: int +) -> list[Array]: + r"""The inverse squared assignment factor :math:`|D_d(m)|^{-2}` of every axis. + + The B-spline spreading multiplies every mode of the charge mesh by + :math:`\prod_d D_d(m_d)`, :math:`D_d(m) = \sum_l M_p(l + 1) + e^{2\pi i m l / n_d}`, which the influence function divides out. + """ + xp = array_api_compat.array_namespace(m_axes[0]) + weights = bspline_integer_weights(order) + factors = [] + for m, n_d in zip(m_axes, dims, strict=True): + theta = (2.0 * math.pi) * m / n_d + real = xp.zeros_like(theta) + imag = xp.zeros_like(theta) + for l, w in enumerate(weights): + real = real + float(w) * xp.cos(theta * l) + imag = imag + float(w) * xp.sin(theta * l) + factors.append(1.0 / (real * real + imag * imag)) + return factors + + +def mesh_energy( + charge: Array, + mesh: MeshGrid, + frame_id: Array, + kernel: BaseInteractionKernel, + order: int, +) -> Array: + r"""Energy of the periodic frames by the smooth particle-mesh Ewald method, per atom. + + The charges spread to the mesh of every frame with B-splines of order + ``order`` along the fractional coordinates; a real FFT gives the half + spectrum :math:`Q(\mathbf{m})` of the charge mesh. The influence function + + .. math:: + G(\mathbf{m}) = \frac{k_e}{V} \hat K(\mathbf{k}_\mathbf{m}) + \prod_d |D_d(m_d)|^{-2}, + \qquad \mathbf{k}_\mathbf{m} = 2\pi B^{-1} \mathbf{m}, + + corrects for the spreading, so that the potential on the mesh is the + inverse transform of :math:`G Q` and the potential at a node its + interpolation by the same B-splines. :math:`|\mathbf{k}|^2` is the + quadratic form :math:`(2\pi)^2 \mathbf{m}^\mathsf{T} W \mathbf{m}` of the + reciprocal metric :math:`W = A^\mathsf{T} A`, :math:`A = S^{-\mathsf{T}} + B^{-1}` of the strained cell, and the strain scales the volume by + :math:`\det S`; the fractional coordinates do not depend on it. The zero + mode carries no weight, nor do the modes at the Nyquist frequency of an + axis, which have no unique wave number. + + Parameters + ---------- + charge + Node charges with shape ``(N, C)``, zero on the nodes that carry none. + mesh + The cells and the mesh of the periodic frames. + frame_id + The frame of every node with shape ``(N,)``. + kernel + The pair kernel. + order + The order of the B-splines. + + Returns + ------- + Array + The energy of every node with shape ``(N,)``, in eV; zero in a frame + masked out by :attr:`MeshGrid.mask`. + """ + xp = array_api_compat.array_namespace(charge) + device = array_api_compat.device(charge) + dtype = charge.dtype + nf = mesh.box.shape[0] + n_node, n_channel = charge.shape + m_full = [xp.astype(m, dtype) for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)] + n1, n2, n3 = (m.shape[0] for m in m_full) + n_mesh = n1 * n2 * n3 + box = xp.astype(mesh.box, dtype) + volume = xp.astype(mesh.volume, dtype) + # === Step 1. Stencils of the nodes on the mesh === + box_inv = _inverse3(xp, box) + # frac_d = sum_j r_j (B^-1)_jd, elementwise rather than a batched product + # of 3x3 matrices. + frac = xp.sum( + xp.astype(mesh.coord, dtype)[:, :, None] * xp.take(box_inv, frame_id, axis=0), + axis=1, + ) + frac = frac - xp.floor(frac) + index, weight = _stencil(xp, frac, (n1, n2, n3), order) + # Flat targets on the frame- and channel-major meshes, (N, C, order**3). + channel = xp.arange(n_channel, dtype=xp.int64, device=device) + first = (frame_id[:, None] * n_channel + channel[None, :]) * n_mesh + target = first[:, :, None] + index[:, None, :] + # === Step 2. Charge mesh and its half spectrum === + q_mesh = xp.zeros((nf * n_channel * n_mesh,), dtype=dtype, device=device) + q_mesh = xp_add_at( + q_mesh, + xp.reshape(target, (-1,)), + xp.reshape(charge[:, :, None] * weight[:, None, :], (-1,)), + ) + q_mesh = xp.reshape(q_mesh, (nf * n_channel, n1, n2, n3)) + spectrum = xp.fft.rfftn(q_mesh, axes=(-3, -2, -1)) + # === Step 3. Influence function on the half spectrum, (nf, n1, n2, h3) === + m3 = m_full[2][: n3 // 2 + 1] + m_half = [m_full[0], m_full[1], m3] + strain = xp.astype(mesh.strain, dtype) + # The rows of A^T = B^-T S^-1 are the reciprocal vectors of the strained + # cell, strained like wave numbers; W = A^T A. + a_t = _strained_rows(xp, xp.matrix_transpose(box_inv), strain) + metric = xp.sum(a_t[:, :, None, :] * a_t[:, None, :, :], axis=-1) + volume = _strained_volume(volume, strain) + g1 = m_half[0][:, None, None] + g2 = m_half[1][None, :, None] + g3 = m_half[2][None, None, :] + # Six monomials of the mode indices times six entries of the metric of + # every frame, elementwise. + monomials = (g1 * g1, g2 * g2, g3 * g3, g1 * g2, g1 * g3, g2 * g3) + entries = ( + metric[:, 0, 0], + metric[:, 1, 1], + metric[:, 2, 2], + 2.0 * metric[:, 0, 1], + 2.0 * metric[:, 0, 2], + 2.0 * metric[:, 1, 2], + ) + k_sq = None + for entry, monomial in zip(entries, monomials, strict=True): + term = entry[:, None, None, None] * monomial[None, ...] + k_sq = term if k_sq is None else k_sq + term + k_sq = (2.0 * math.pi) ** 2 * k_sq + # The zero mode and the Nyquist planes carry no weight. + keep = ( + (2.0 * xp.abs(g1) != n1) & (2.0 * xp.abs(g2) != n2) & (2.0 * xp.abs(g3) != n3) + ) + keep = keep & ((g1 != 0.0) | (g2 != 0.0) | (g3 != 0.0)) + keep = keep[None, ...] & mesh.mask[:, None, None, None] + safe_k_sq = xp.where(keep, k_sq, xp.ones_like(k_sq)) + shape = safe_k_sq.shape + # One row per frame and first mode index, along which the gradients of + # the kernel parameters sum first. + rows = xp.reshape(safe_k_sq, (nf * n1, n2 * (n3 // 2 + 1))) + transform = xp.reshape(kernel.kspace(rows), shape) + f1, f2, f3 = _assignment_factor(m_half, (n1, n2, n3), order) + factor = f1[:, None, None] * f2[None, :, None] * f3[None, None, :] + influence = (COULOMB_CONSTANT * transform * factor[None, ...]) / volume[ + :, None, None, None + ] + influence = xp.where(keep, influence, xp.zeros_like(influence)) + # === Step 4. Potential on the mesh and at the nodes === + spectrum = xp.reshape(spectrum, (nf, n_channel, n1, n2, n3 // 2 + 1)) + potential_mesh = xp.fft.irfftn( + spectrum * influence[:, None, ...], + s=(n1, n2, n3), + axes=(-3, -2, -1), + norm="forward", + ) + stencil_potential = xp.take( + xp.reshape(potential_mesh, (-1,)), xp.reshape(target, (-1,)), axis=0 + ) + potential = xp.sum( + xp.reshape(stencil_potential, (n_node, n_channel, -1)) * weight[:, None, :], + axis=-1, + ) + periodic_node = xp.take(mesh.mask, frame_id, axis=0) + potential = potential - _self_potential(kernel, charge, periodic_node) + return 0.5 * xp.sum(charge * potential, axis=-1) diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 7f3f8b9705..b5549d4adb 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -73,3 +73,5 @@ def compute_fitting_input_stats( The lazy sampled function to get data frames from different data systems. """ self.get_fitting_net().compute_input_stats(sampled_func) + if self.atomic_model.has_long_range(): + self.atomic_model.long_range.compute_input_stats(sampled_func) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 607f7dfb83..4788e27b29 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -48,6 +48,7 @@ NeighborGraph, build_neighbor_graph, build_neighbor_graph_ase, + cast_topology_geometry, compact_nodes, expand_node_values, ) @@ -564,12 +565,20 @@ def _call_common_graph( # atom count the mask is all true and this is a renumbering by the # identity. ng, node_index = compact_nodes(ng, atype_flat >= 0) + atype_real = xp.take(atype_flat, node_index, axis=0) + topologies = self.atomic_model.build_topologies( + xp.take(xp.reshape(cc, (n_padded, 3)), node_index, axis=0), + ng.n_node, + bb, + atype=atype_real, + pair_excl=pair_excl, + ) # OUTPUT-AGNOSTIC standard model dict (````, ``_redu``, # derivative name-holders ``None``, plus int ``mask``), like the # dense ``call_common``. ``call_lower_graph`` masks virtual atoms # (atype<0) and sets the real int mask. model_predict = self.call_lower_graph( - atype=xp.take(atype_flat, node_index, axis=0), + atype=atype_real, n_node=ng.n_node, edge_index=ng.edge_index, edge_vec=ng.edge_vec, @@ -589,6 +598,7 @@ def _call_common_graph( else None ), charge_spin=charge_spin, + topologies=topologies, ) # Public ABI is rectangular (nf, nloc, *); the lower is flat over # the real atoms. Scatter per-atom keys back onto the padded width @@ -731,6 +741,7 @@ def forward_common_atomic_graph( comm_dict: dict | None = None, charge_spin: Array | None = None, spin: Array | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, Array]: """Model-level graph forward (no type cast). Analogue of the dense :meth:`forward_common_atomic`. @@ -773,6 +784,9 @@ def forward_common_atomic_graph( Per-node spin vectors, flat (N, 3), or ``None``. Forwarded unchanged to the atomic model's ``forward_common_atomic_graph`` (and, from there, the descriptor's ``call_graph``). + topologies + The topologies the atomic model declares by + ``required_topologies()``, on the node axis of the graph. Returns ------- @@ -795,6 +809,7 @@ def forward_common_atomic_graph( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) return fit_output_to_model_output_graph( atomic_ret, @@ -817,6 +832,7 @@ def call_common_lower_graph( comm_dict: dict | None = None, charge_spin: Array | None = None, spin: Array | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, Array]: """Graph-native PUBLIC lower (dpa1/se_atten concat-tebd, attention included). @@ -858,6 +874,9 @@ def call_common_lower_graph( Per-node spin vectors, flat (N, 3), or ``None``. Cast to the model precision alongside the other node inputs and forwarded unchanged to :meth:`forward_common_atomic_graph`. + topologies + The topologies the atomic model declares by + ``required_topologies()``, on the node axis of the graph. Returns ------- @@ -871,6 +890,12 @@ def call_common_lower_graph( charge_spin=charge_spin, spin=spin, ) + if topologies: + # The topologies carry geometry like the edge vectors. + topologies = { + name: cast_topology_geometry(topology, edge_vec.dtype) + for name, topology in topologies.items() + } model_predict = self.forward_common_atomic_graph( atype, n_node, @@ -883,6 +908,7 @@ def call_common_lower_graph( comm_dict=comm_dict, charge_spin=cs, spin=sp, + topologies=topologies, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 66d32f8f8b..21c3be799b 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -13,6 +13,9 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) +from deepmd.dpmodel.long_range import ( + BaseLongRange, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -47,6 +50,7 @@ atomic_model=DPAtomicModel, pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, + long_range_base=BaseLongRange, ) get_zbl_model = _model_factory.get_zbl_model diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 38a4f4b9e4..118b09a75f 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -65,6 +65,83 @@ def get_model_components( return descriptor, fitting, fitting_type +def get_long_range( + data: dict, + descriptor: Any, + fitting: Any, + *, + long_range_base: type | None, + backend_name: str, +) -> Any | None: + """Construct the long-range module a model configuration adds to its fitting. + + The module reads the descriptor of the model, and its networks take the + conditioning inputs of the energy fitting they accompany. + + Parameters + ---------- + data : dict + The model configuration, with the optional ``long_range`` section. + descriptor : Any + The descriptor of the model. + fitting : Any + The energy fitting of the model. + long_range_base : type or None + Backend long-range registry base class; ``None`` for a backend + without long-range modules. + backend_name : str + Backend name used in error messages. + + Returns + ------- + Any or None + The backend-native long-range module, or ``None`` when none is + configured. + + Raises + ------ + NotImplementedError + If a long-range module is configured for a backend without them. + """ + config = data.get("long_range") + if config is None: + return None + if long_range_base is None: + raise NotImplementedError( + f"long-range modules are not supported in the {backend_name} backend" + ) + config = copy.deepcopy(config) + long_range_cls = long_range_base.get_class_by_type(config.pop("type")) + return long_range_cls.from_config( + config, + descriptor, + fitting, + fitting_seed=data.get("fitting_net", {}).get("seed"), + ) + + +def _long_range_kwargs( + data: dict, + descriptor: Any, + fitting: Any, + *, + long_range_base: type | None, + backend_name: str, +) -> dict[str, Any]: + """The ``long_range`` constructor argument, absent when none is configured. + + Model classes without a long-range module keep their constructor contract. + """ + long_range = get_long_range( + data, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ) + return {} if long_range is None else {"long_range": long_range} + + def get_standard_model( data: dict, *, @@ -72,6 +149,7 @@ def get_standard_model( fitting_base: type, model_base: type, backend_name: str, + long_range_base: type | None = None, ) -> Any: """Construct a standard model through backend registries.""" descriptor, fitting, fitting_type = get_model_components( @@ -89,6 +167,13 @@ def get_standard_model( atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=data.get("pair_exclude_types", []), preset_out_bias=data.get("preset_out_bias"), + **_long_range_kwargs( + data, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ), ) @@ -103,6 +188,11 @@ def get_zbl_model( backend_name: str, ) -> Any: """Construct a ZBL model from backend-native atomic model classes.""" + if data.get("long_range") is not None: + raise NotImplementedError( + "a long-range module is not supported with the tabulated ZBL " + "potential (`use_srtab`)" + ) descriptor, fitting, fitting_type = get_model_components( data, descriptor_base=descriptor_base, @@ -141,6 +231,7 @@ def get_linear_atomic_model( pairtab_model: type, linear_atomic_model: type | None = None, descriptor_child_builder: "Callable[[dict], Any | None] | None" = None, + long_range_base: type | None = None, ) -> Any: """Build the ``LinearEnergyAtomicModel`` composition from a config. @@ -181,6 +272,9 @@ def get_linear_atomic_model( back to the generic registry build. Backends use it to route family-specific model types (e.g. DPA4/SeZM) through their validated builders. + long_range_base : type, optional + Backend long-range registry base class, which builds the long-range + module of a child; ``None`` for a backend without long-range modules. Raises ------ @@ -188,8 +282,11 @@ def get_linear_atomic_model( If more than one ``inner_potential`` child is given, if an ``inner_potential`` child has no unique learned sibling, if a bridged composition does not use ``weights: "sum"``, if a child - carries a ``bridging_method`` flag, or if a child is of an + carries a ``bridging_method`` flag, if a child with a long-range + module carries its own ``pair_exclude_types``, or if a child is of an unsupported kind. + NotImplementedError + If a child with a long-range module carries ``atom_exclude_types``. """ from deepmd.dpmodel.atomic_model.inner_potential import ( InnerPotentialAtomicModel, @@ -311,6 +408,24 @@ def get_linear_atomic_model( if descriptor_child_builder is not None: child = descriptor_child_builder(sub) if child is None: + # A long-range child is checked as the same model alone: the + # pair exclusion of the composition shapes the graph and the + # pair topology the children share, and atom exclusion is + # unsupported with a long-range module. + if sub.get("long_range") is not None: + if sub.get("pair_exclude_types"): + raise ValueError( + "a linear_ener sub-model with a long-range module " + "takes `pair_exclude_types` from the composition, " + "whose sub-models share one neighbor graph and one " + "pair topology; set it on the composition" + ) + if sub.get("atom_exclude_types"): + raise NotImplementedError( + "`atom_exclude_types` is not supported with a " + "long-range module, which attributes each pair " + "energy to both atoms of the pair" + ) descriptor, fitting, _ = get_model_components( sub, descriptor_base=descriptor_base, @@ -325,6 +440,13 @@ def get_linear_atomic_model( fitting, type_map=sub["type_map"], preset_out_bias=preset_out_bias, + **_long_range_kwargs( + sub, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ), ) built[i] = child else: @@ -467,6 +589,7 @@ def __init__( pairtab_model: type | None = None, zbl_model: type | None = None, linear_atomic_model: type | None = None, + long_range_base: type | None = None, ) -> None: """Store backend-native classes used by all model construction paths.""" self.descriptor_base = descriptor_base @@ -477,6 +600,7 @@ def __init__( self.pairtab_model = pairtab_model self.zbl_model = zbl_model self.linear_atomic_model = linear_atomic_model + self.long_range_base = long_range_base def get_model_components(self, data: dict) -> tuple[Any, Any, str]: """Construct descriptor and fitting objects for this backend.""" @@ -495,6 +619,17 @@ def get_standard_model(self, data: dict) -> Any: fitting_base=self.fitting_base, model_base=self.model_base, backend_name=self.backend_name, + long_range_base=self.long_range_base, + ) + + def get_long_range(self, data: dict, descriptor: Any, fitting: Any) -> Any | None: + """Construct the long-range module of a model config for this backend.""" + return get_long_range( + data, + descriptor, + fitting, + long_range_base=self.long_range_base, + backend_name=self.backend_name, ) def get_linear_atomic_model( @@ -515,6 +650,7 @@ def get_linear_atomic_model( pairtab_model=self.pairtab_model, linear_atomic_model=self.linear_atomic_model, descriptor_child_builder=descriptor_child_builder, + long_range_base=self.long_range_base, ) def get_zbl_model(self, data: dict) -> Any: diff --git a/deepmd/dpmodel/model/native_spin_model.py b/deepmd/dpmodel/model/native_spin_model.py index 5df83d91df..6fdd5744d6 100644 --- a/deepmd/dpmodel/model/native_spin_model.py +++ b/deepmd/dpmodel/model/native_spin_model.py @@ -79,6 +79,10 @@ class NSM(T_Model, NativeSpinModelKind): def __init__(self, *args: Any, spin: Spin, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + if self.atomic_model.has_long_range(): + raise NotImplementedError( + "native spin is not supported for a model with a long-range module" + ) self.spin = spin self.ntypes_real = self.spin.ntypes_real # Per-real-type 0/1 spin gate, derived from ``use_spin`` and hence diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index ae7ed8bae3..c1bb8c3fe1 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -122,6 +122,11 @@ def __init__( spin: Spin, ) -> None: super().__init__() + if backbone_model.atomic_model.has_long_range(): + raise NotImplementedError( + "the virtual-atom spin scheme is not supported for a model with " + "a long-range module" + ) self.backbone_model = backbone_model self.spin = spin # Spin graph-lower unsupported: carry-all routing diverges on diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 6895ce3e1d..f13d80cb46 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -5,8 +5,12 @@ ``graph`` (the ``NeighborGraph``/``GraphLayout`` contract + derived node-validity + 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). +``from_dense_quartet`` legacy converter), ``pair_graph`` (the complete +intra-frame pair topology of interactions without a cutoff), +``reciprocal_graph`` (the reciprocal-space topologies of periodic frames), +``topology`` (building and laying out the topologies a model declares), +``segment`` (mask-aware segment-reduction toolkit), and ``derivatives`` (edge +and pair force/virial assembly). See the design discussion wanghan-iapcm/deepmd-kit#4. """ @@ -33,6 +37,7 @@ ) from .derivatives import ( edge_force_virial, + pair_force_virial, ) from .env import ( edge_env_mat, @@ -53,19 +58,56 @@ pad_and_guard_angles, pad_and_guard_edges, ) +from .pair_graph import ( + ALL_PAIRS, + PairGraph, + build_pair_graph, + check_complete_frames, + unordered_pair_index, +) from .pairs import ( center_edge_pairs, ) +from .reciprocal_graph import ( + K_BLOCK, + K_GRID, + PME_MESH, + KGrid, + MeshGrid, + build_fixed_k_grid, + build_k_grid, + build_mesh_grid, + check_periodic_cells, + enumerate_k_grid_bounds, + k_grid_bounds_of_cells, + mesh_dims_of_cells, +) from .segment import ( segment_max, segment_mean, segment_softmax, segment_sum, ) +from .topology import ( + TOPOLOGY_TYPES, + build_topologies, + cast_topology_geometry, + flatten_topologies, + topology_flat_fields, + unflatten_topologies, +) __all__ = [ + "ALL_PAIRS", + "K_BLOCK", + "K_GRID", + "PME_MESH", + "TOPOLOGY_TYPES", "GraphLayout", + "KGrid", + "MeshGrid", "NeighborGraph", + "PairGraph", "angle_padding_fraction", "angle_to_edge_sum", "angle_to_node_sum", @@ -75,25 +117,41 @@ "attach_edge_csr", "build_angle_index", "build_edge_csr", + "build_fixed_k_grid", + "build_k_grid", + "build_mesh_grid", "build_neighbor_graph", "build_neighbor_graph_ase", + "build_pair_graph", + "build_topologies", "canonicalize_neighbor_graph", + "cast_topology_geometry", "center_edge_pairs", + "check_complete_frames", + "check_periodic_cells", "compact_nodes", "edge_env_mat", "edge_force_virial", + "enumerate_k_grid_bounds", "expand_node_values", + "flatten_topologies", "frame_id_from_n_node", "from_dense_quartet", "graph_angle_cos", "graph_from_dense_quartet", + "k_grid_bounds_of_cells", + "mesh_dims_of_cells", "neighbor_graph_from_ijs", "node_ownership_mask", "node_validity_mask", "pad_and_guard_angles", "pad_and_guard_edges", + "pair_force_virial", "segment_max", "segment_mean", "segment_softmax", "segment_sum", + "topology_flat_fields", + "unflatten_topologies", + "unordered_pair_index", ] diff --git a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py index 09111fa5d5..5e3f34ec04 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -118,3 +118,71 @@ def edge_force_virial( 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 + + +def pair_force_virial( + g_p: Array, + pair_vec: Array, + pair_index: Array, + pair_mask: Array, + n_node: Array, + node_capacity: int | None = None, +) -> tuple[Array, Array, Array]: + """Assemble per-node force/atom-virial and PER-FRAME virial of a pair set. + + A pair set lists every unordered pair once (see + :class:`~deepmd.dpmodel.utils.neighbor_graph.pair_graph.PairGraph`), so its + per-pair gradient ``g_p = dE/d(pair_vec)`` acts on both endpoints: + ``pair_vec = r_src - r_dst`` gives ``F_dst += g_p`` and ``F_src -= g_p``, + as for an edge. The pair virial ``w = -g_p (x) pair_vec`` is split evenly + between the two endpoints, which keeps the atomic virial independent of + the order in which a pair is listed; its per-frame sum is the sum of + ``w`` over the pairs of that frame. + + Parameters + ---------- + g_p + (P, 3) per-pair gradient ``dE/d(pair_vec)``. + pair_vec + (P, 3) per-pair displacement ``r_src - r_dst``. + pair_index + (2, P) ``[src, dst]`` node endpoints of each pair. + pair_mask + (P,) boolean real-pair mask; guard pairs (``False``) are zeroed + before any scatter. + n_node + (nf,) per-frame node counts. + node_capacity + Size of the node axis ``N``; ``None`` takes ``sum(n_node)``. + + Returns + ------- + force + (N, 3) per-node force. + atom_virial + (N, 3, 3) per-node virial, half of each pair virial per endpoint. + virial + (nf, 3, 3) per-frame virial. + """ + xp = array_api_compat.array_namespace(g_p) + n_out = node_capacity if node_capacity is not None else int(xp.sum(n_node)) + nf = n_node.shape[0] + if isinstance(n_out, int) and n_out == 0: + device = array_api_compat.device(g_p) + return ( + xp.zeros((0, 3), dtype=g_p.dtype, device=device), + xp.zeros((0, 3, 3), dtype=g_p.dtype, device=device), + xp.zeros((nf, 3, 3), dtype=g_p.dtype, device=device), + ) + g = g_p * xp.astype(pair_mask[:, None], g_p.dtype) + # As for edges, the modulo bounds the guard endpoints under export. + src = pair_index[0] % n_out + dst = pair_index[1] % n_out + force = segment_sum(g, dst, n_out) - segment_sum(g, src, n_out) + w_half = -0.5 * (g[:, :, None] * pair_vec[:, None, :]) # (P, 3, 3) + atom_virial = segment_sum(w_half, src, n_out) + segment_sum(w_half, dst, n_out) + # Both endpoints of a pair share a frame, so reducing the per-node virial + # over nodes equals reducing the pair virial over pairs. + node_frame = frame_id_from_n_node(n_node, n_total=n_out) + virial = segment_sum(atom_virial, node_frame, nf) + return force, atom_virial, virial diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py new file mode 100644 index 0000000000..e474abd4ed --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Complete intra-frame pair topology for interactions without a cutoff. + +A module whose interaction acts between every two atoms of a system, such as +a non-periodic long-range module, cannot read its geometry from the +short-range :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph`. It +declares the :data:`ALL_PAIRS` topology instead: every unordered pair of valid +nodes sharing a frame, listed once with ``src < dst``. The pairs are generated +from the per-frame node counts alone, so their storage is the true pair count +``sum_f n_f (n_f - 1) / 2`` with neither a distance search nor padding to the +widest frame. Pairs of types excluded by the model-level pair exclusion are +left out, as they are left out of the neighbor graph. Like the edge axis of +the neighbor graph, the pair axis ends in masked guard pairs, which keep its +length a dynamic dimension of its own however few real pairs a batch holds. + +Geometry enters only through :attr:`PairGraph.vec`, which an autograd backend +makes a leaf next to the short-range ``edge_vec``; its gradient is pulled back +to the nodes by :func:`~deepmd.dpmodel.utils.neighbor_graph.derivatives.pair_force_virial`. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + ClassVar, +) + +import array_api_compat + +from .graph import ( + frame_id_from_n_node, + pad_and_guard_edges, +) +from .segment import ( + segment_sum, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + +ALL_PAIRS = "all_pairs" +"""Name of the topology holding every unordered pair of valid nodes of a frame.""" + + +@dataclass +class PairGraph: + """Unordered pairs of nodes on the flat node axis of a NeighborGraph.""" + + index: Array + """(2, P) int64 [src, dst] with src < dst; both endpoints in one frame.""" + vec: Array + """(P, 3) float r_src - r_dst; the geometry leaf of the pair interaction.""" + mask: Array + """(P,) bool real (1) vs guard (0) pairs; guards join node 0 to itself.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ("vec",) + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("vec",) + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "index": {1: "entry"}, + "vec": {0: "entry"}, + "mask": {0: "entry"}, + } + """Dynamic axes of every field: the pair axis.""" + + +def unordered_pair_index(n_node: Array, node_mask: Array | None = None) -> Array: + """Enumerate the unordered pairs of valid nodes within every frame. + + Within a frame of ``n`` valid nodes the pairs follow the row-major order + of the strict upper triangle, ``(0, 1), (0, 2), ..., (0, n-1), (1, 2), ...``. + Row ``i`` starts at rank ``S(i) = i (2n - 1 - i) / 2``; the row of rank + ``k`` is the floor of the smaller root of ``S(i) = k``, evaluated in + float64 and corrected by one integer step against rounding. + + The pair count is data dependent, so the enumeration runs eagerly. + + Parameters + ---------- + n_node + Per-frame node counts with shape ``(nf,)``. + node_mask + Validity of the flat nodes with shape ``(N,)``; ``None`` takes every + node as valid. Invalid nodes appear in no pair. + + Returns + ------- + Array + The ``[src, dst]`` node indices with shape ``(2, P)``, int64. + """ + xp = array_api_compat.array_namespace(n_node) + device = array_api_compat.device(n_node) + n_node = xp.astype(n_node, xp.int64) + if node_mask is None: + count = n_node + else: + frame_id = frame_id_from_n_node(n_node, n_total=node_mask.shape[0]) + count = segment_sum(xp.astype(node_mask, xp.int64), frame_id, n_node.shape[0]) + n_pair = count * (count - 1) // 2 + total = int(xp.sum(n_pair)) + if total == 0: + return xp.zeros((2, 0), dtype=xp.int64, device=device) + pair_end = xp.cumulative_sum(n_pair) + rank = xp.arange(total, dtype=xp.int64, device=device) + pair_frame = xp.astype(xp.searchsorted(pair_end, rank, side="right"), xp.int64) + local = rank - xp.take(pair_end - n_pair, pair_frame, axis=0) + n = xp.take(count, pair_frame, axis=0) + + def row_start(row: Array) -> Array: + return row * (2 * n - 1 - row) // 2 + + b = xp.astype(2 * n - 1, xp.float64) + discriminant = b * b - 8.0 * xp.astype(local, xp.float64) + discriminant = xp.where( + discriminant > 0.0, discriminant, xp.zeros_like(discriminant) + ) + row = xp.astype(xp.floor((b - xp.sqrt(discriminant)) / 2.0), xp.int64) + row = xp.where(local < row_start(row), row - 1, row) + row = xp.where(local >= row_start(row + 1), row + 1, row) + col = row + 1 + (local - row_start(row)) + first = xp.take(xp.cumulative_sum(count) - count, pair_frame, axis=0) + src = first + row + dst = first + col + if node_mask is not None: + # The ranks above count valid nodes only; map them to node positions. + position = xp.reshape(xp.nonzero(node_mask)[0], (-1,)) + position = xp.astype(position, xp.int64) + src = xp.take(position, src, axis=0) + dst = xp.take(position, dst, axis=0) + return xp.stack([src, dst]) + + +def build_pair_graph( + coord: Array, + n_node: Array, + atype: Array, + node_mask: Array | None = None, + pair_excl: PairExcludeMask | None = None, +) -> PairGraph: + """Build the :data:`ALL_PAIRS` topology of flat node coordinates. + + Parameters + ---------- + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + A differentiable input keeps :attr:`PairGraph.vec` differentiable. + n_node + Per-frame node counts with shape ``(nf,)``. + atype + Node types with shape ``(N,)``; invalid nodes carry any type. + node_mask + Validity of the nodes with shape ``(N,)``, or ``None`` when every + node is valid. + pair_excl + The model-level pair exclusion, or ``None``; pairs of excluded types + are left out. + + Returns + ------- + PairGraph + The unordered pairs of valid nodes and their displacements, followed + by two masked guard pairs. + """ + xp = array_api_compat.array_namespace(coord) + index = unordered_pair_index(n_node, node_mask) + if pair_excl is not None: + keep = pair_excl.build_edge_exclude_mask(index, atype) + index = xp.take(index, xp.reshape(xp.nonzero(keep)[0], (-1,)), axis=1) + vec = xp.take(coord, index[0], axis=0) - xp.take(coord, index[1], axis=0) + index, vec, mask = pad_and_guard_edges(index, vec, None) + return PairGraph(index=index, vec=vec, mask=mask) + + +# Guard pairs of the pair topology of a batch without non-periodic frames. +_GUARD_PAIRS = 4 + + +def empty_pair_graph(like: Array) -> PairGraph: + """The pair topology of a batch without non-periodic frames: guard pairs only. + + The guard pairs number :data:`_GUARD_PAIRS`, unlike the two rows of the + index and the three components of the vectors: a compiled graph that met + a pair axis of the length of one of those would identify the two, and + split at the reductions over the pairs. + """ + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + index = xp.zeros((2, 0), dtype=xp.int64, device=device) + vec = xp.zeros((0, 3), dtype=like.dtype, device=device) + index, vec, mask = pad_and_guard_edges(index, vec, None, min_edges=_GUARD_PAIRS) + return PairGraph(index=index, vec=vec, mask=mask) + + +def check_complete_frames(box: Array | None) -> None: + """Reject a cell, since the all-pairs topology covers non-periodic frames. + + Parameters + ---------- + box + Simulation cells, or ``None`` for non-periodic systems. + + Raises + ------ + NotImplementedError + If a cell is given. + """ + if box is not None: + raise NotImplementedError( + "the all-pairs topology of the long-range modules describes " + "complete non-periodic systems only; a module evaluates periodic " + "cells with `periodic: true` (pass box=None for a non-periodic " + "system)" + ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py new file mode 100644 index 0000000000..c337cf40f4 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Reciprocal-space topologies of the periodic frames of a batch. + +A long-range module that serves periodic systems evaluates the lattice sum of +its pair kernel in reciprocal space, which needs the cell of every frame and +the absolute node positions rather than pair displacements. It declares one of +two topologies beside +:data:`~deepmd.dpmodel.utils.neighbor_graph.pair_graph.ALL_PAIRS`: + +- :data:`K_GRID` lists the wave numbers of every frame explicitly, ragged over + the frames, for a direct sum over them; +- :data:`PME_MESH` carries the cells and the mode indices of a real-space mesh + shared by the frames, for a smooth particle-mesh evaluation. + +Both carry the node positions and a zero strain as geometry: an autograd +backend makes them leaves, whose gradients are the force and the negative +virial of the reciprocal-space energy. A batch without cells comes back with +masked topologies, so a module serving periodic systems also serves +non-periodic batches through its pair topology. +""" + +from __future__ import ( + annotations, +) + +import math +import warnings +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + ClassVar, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + +K_GRID = "k_grid" +"""Name of the topology holding the reciprocal lattice points of periodic frames.""" + +PME_MESH = "pme_mesh" +"""Name of the topology carrying the cells and the mesh of periodic frames.""" + + +@dataclass +class KGrid: + r"""Reciprocal lattice points of the periodic frames of a batch, frame by frame. + + Row ``f`` of :attr:`vec` holds the wave numbers :math:`\mathbf{k} = 2\pi + B_f^{-1}\mathbf{n}` of frame ``f``, for integer :math:`\mathbf{n}`, built + eagerly outside any compiled region, in blocks of :data:`K_BLOCK` points + padded with masked guard points to the blocks of the longest row; row + ``f`` of :attr:`slot` holds the nodes of frame ``f``, padded with ``-1``. + A frame thus pairs its own wave numbers with its own nodes alone, at a + cost of the frame count times the padded wave-number count and the + padded node count, and every sum over the wave numbers of a frame runs + over the points of a block first. The block, point and slot axes are dynamic dimensions of + their own, of at least two entries. + + The energy is evaluated on the strained geometry of :attr:`strain`: + positions by ``S = I + strain``, wave numbers by :math:`S^{-T}` and volumes + by :math:`\det S`, so that the gradient with respect to the strain is the + negative virial; the phases :math:`\mathbf{k} \cdot \mathbf{r}` do not + change under the strain. The strain is zero, and the evaluation takes it + to first order, which fixes the energy and that gradient. The graph lower + makes :attr:`coord` and :attr:`strain` autograd leaves at its inputs. + """ + + vec: Array + """(nf, nb, kb, 3) float the wave numbers of every frame in blocks, in Å⁻¹.""" + mask: Array + """(nf, nb, kb) bool real (1) vs guard (0) points.""" + slot: Array + """(nf, ns) int64 the nodes of every frame, ``-1`` past its last node.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array + """(N, 3) float node positions, in Å.""" + strain: Array + """(nf, 3, 3) float zero strain of every frame.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ("vec", "volume", "coord", "strain") + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("coord", "strain") + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "vec": {0: "frame", 1: "entry", 2: "k_point"}, + "mask": {0: "frame", 1: "entry", 2: "k_point"}, + "slot": {0: "frame", 1: "slot"}, + "volume": {0: "frame"}, + "coord": {0: "node"}, + "strain": {0: "frame"}, + } + """Dynamic axes of every field: the frame, block, point, slot and node axes.""" + + +@dataclass +class MeshGrid: + r"""Cells and mesh of the periodic frames of a batch. + + The frames share a real-space mesh of :math:`n_1 \times n_2 \times n_3` + points spanning each cell, whose integer mode indices along every axis + are the fields :attr:`mode_1`, :attr:`mode_2` and :attr:`mode_3`: their + lengths are the mesh dimensions, so a compiled or exported graph takes the + mesh as dynamic dimensions. Every dimension is even. A batch without a + cell comes back with identity cells, a minimal mesh and an all-false + :attr:`mask`. + + The geometry contract is that of :class:`KGrid`: the energy is evaluated + with the positions strained by ``S = I + strain`` and the cells by + :math:`B S^{\mathsf{T}}`, the lattice vectors being the rows of + :math:`B`, so that the fractional coordinates are invariant and the + gradient with respect to the strain is the negative virial. The volumes + of the unstrained cells are computed eagerly, and the evaluation scales + them by :math:`\det S`, to first order in the strain. + """ + + box: Array + """(nf, 3, 3) float simulation cells; rows are the lattice vectors, in Å.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array + """(N, 3) float node positions, in Å.""" + strain: Array + """(nf, 3, 3) float zero strain of every frame.""" + mask: Array + """(nf,) bool periodic (1) vs non-periodic (0) frames.""" + mode_1: Array + """(n1,) float mode indices of the first mesh axis, in FFT order.""" + mode_2: Array + """(n2,) float mode indices of the second mesh axis, in FFT order.""" + mode_3: Array + """(n3,) float mode indices of the third mesh axis, in FFT order.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ( + "box", + "volume", + "coord", + "strain", + "mode_1", + "mode_2", + "mode_3", + ) + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("coord", "strain") + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "box": {0: "frame"}, + "volume": {0: "frame"}, + "coord": {0: "node"}, + "strain": {0: "frame"}, + "mask": {0: "frame"}, + "mode_1": {0: "mesh_1"}, + "mode_2": {0: "mesh_2"}, + "mode_3": {0: "mesh_3"}, + } + """Dynamic axes of every field: the frame axis, the node axis, the mesh axes.""" + + +def check_periodic_cells(box: Array) -> None: + """Reject a cell that spans no lattice. + + A cell whose volume is a vanishing fraction of the product of the lengths + of its lattice vectors, or that is not finite, spans no lattice; an + all-zero cell is no marker of a non-periodic frame, which takes no cell. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 9)`` or ``(nf, 3, 3)``; rows are + the lattice vectors. + + Raises + ------ + ValueError + If a cell spans no lattice. + """ + cells = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + volume = np.abs(np.linalg.det(cells)) + spanned = volume > 1e-6 * np.prod(np.linalg.norm(cells, axis=-1), axis=-1) + if not np.all(spanned): + f = int(np.argmin(spanned)) + raise ValueError( + f"the cell of periodic frame {f} has no volume; a periodic frame " + "needs three independent lattice vectors, and a non-periodic " + "system takes no cell" + ) + + +def _cells(box: Array) -> Array: + """The cells with shape ``(nf, 3, 3)``, from a ``(nf, 9)`` or ``(nf, 3, 3)`` array.""" + xp = array_api_compat.array_namespace(box) + return xp.reshape(box, (box.shape[0], 3, 3)) + + +def _volume(box: Array) -> Array: + """The cell volumes |a_1 . (a_2 x a_3)| with shape ``(nf,)``.""" + xp = array_api_compat.array_namespace(box) + triple = box[:, 0, :] * xp.linalg.cross(box[:, 1, :], box[:, 2, :]) + return xp.abs(xp.sum(triple, axis=-1)) + + +def _zero_strain(box: Array) -> Array: + xp = array_api_compat.array_namespace(box) + return xp.zeros( + (box.shape[0], 3, 3), dtype=box.dtype, device=array_api_compat.device(box) + ) + + +def k_grid_bounds_of_cells(box: Array, k_max: float) -> np.ndarray: + r"""Per-axis integer wave-number bounds the cells need at ``k_max``. + + A wave number :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` with + :math:`|\mathbf{k}| \le k_{max}` has :math:`|n_d| \le k_{max} |a_d| / 2\pi` + along every lattice vector :math:`\mathbf{a}_d` (Cauchy-Schwarz), so the + bound of every frame is the ceiling of the right-hand side. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + k_max + The wave-number cutoff, in Å⁻¹. + + Returns + ------- + np.ndarray + The per-frame, per-axis integer bounds with shape ``(nf, 3)``. + """ + cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + length = np.linalg.norm(cell, axis=-1) + return np.ceil(k_max * length / (2.0 * math.pi)).astype(np.int64) + + +def enumerate_k_grid_bounds(bounds: Array | list[int]) -> np.ndarray: + r"""Enumerate the integer wave-number indices within per-axis bounds. + + The grid holds every :math:`\mathbf{n} \neq \mathbf{0}` with + :math:`|n_d| \le n_d^{max}`, in lexicographic order with the first axis + outermost. The zero vector is left out: the transform of a Coulomb-like + kernel is singular there, and the tin-foil convention drops that mode. + + Parameters + ---------- + bounds + The per-axis bounds :math:`n_d^{max}` with shape ``(3,)``. + + Returns + ------- + np.ndarray + The integer indices with shape ``(nk, 3)``, int64. + + Raises + ------ + ValueError + If the bounds are not three non-negative integers. + """ + bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(-1) + if bounds.shape != (3,) or np.any(bounds < 0): + raise ValueError( + f"k-grid bounds must be three non-negative integers, got {bounds}" + ) + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=np.int64) for n in bounds], indexing="ij" + ) + n_grid = np.stack(grids, axis=-1).reshape(-1, 3) + return n_grid[np.any(n_grid != 0, axis=-1)] + + +K_BLOCK = 128 +"""Wave numbers per block of a row of a :class:`KGrid`.""" +# A guard point carries a nonzero wave number, whose kernel transform stays +# finite, and no weight. +_GUARD_K = (1e-6, 0.0, 0.0) +# The fewest entries of the block and slot axes, which keeps each a dynamic +# dimension of its own however few points or nodes a batch holds. +_MIN_ENTRIES = 2 + + +def _k_blocks(rows: list[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + """Lay the wave numbers of every frame out in blocks, padded with guard points. + + Returns the wave numbers with shape ``(nf, nb, K_BLOCK, 3)`` and their + mask with shape ``(nf, nb, K_BLOCK)``. + """ + longest = max([row.shape[0] for row in rows], default=0) + blocks = max(-(-longest // K_BLOCK), _MIN_ENTRIES) + vec = np.tile(np.asarray(_GUARD_K), (len(rows), blocks * K_BLOCK, 1)) + mask = np.zeros((len(rows), blocks * K_BLOCK), dtype=bool) + for f, row in enumerate(rows): + vec[f, : row.shape[0]] = row + mask[f, : row.shape[0]] = True + return ( + vec.reshape(len(rows), blocks, K_BLOCK, 3), + mask.reshape(len(rows), blocks, K_BLOCK), + ) + + +def _frame_slots(n_node: Array) -> np.ndarray: + """The nodes of every frame, one row each, padded with ``-1`` to the longest frame.""" + counts = np.asarray(to_numpy_array(n_node), dtype=np.int64).reshape(-1) + width = max(int(counts.max(initial=0)), _MIN_ENTRIES) + first = np.cumsum(counts) - counts + column = np.arange(width, dtype=np.int64)[None, :] + return np.where(column < counts[:, None], first[:, None] + column, -1) + + +def _k_grid( + box: Array, vec: np.ndarray, mask: np.ndarray, slot: np.ndarray, coord: Array +) -> KGrid: + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + return KGrid( + vec=xp.astype(xp.asarray(vec, device=device), box.dtype), + mask=xp.asarray(mask, device=device), + slot=xp.asarray(slot, device=device), + volume=xp.astype(_volume(box), box.dtype), + coord=coord, + strain=_zero_strain(box), + ) + + +def build_k_grid(box: Array, k_max: float, coord: Array, n_node: Array) -> KGrid: + r"""Enumerate the reciprocal lattice points of the periodic frames. + + With the lattice vectors as the rows of the cell :math:`B`, the wave + numbers of a frame are :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` over + integer :math:`\mathbf{n}` with :math:`0 < |\mathbf{k}| \le k_{max}`, so + that :math:`\mathbf{a}_d \cdot \mathbf{k} = 2\pi n_d`. Their count is data + dependent, so the enumeration runs eagerly. The points of every frame + follow the lexicographic order of their integer indices, and the rows of + the frames end in masked guard points up to the blocks of the longest + row. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + k_max + The wave-number cutoff, in Å⁻¹. + coord + Node positions with shape ``(N, 3)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + + Returns + ------- + KGrid + The reciprocal lattice points of every frame, one row each. + """ + box = _cells(box) + cells = np.asarray(to_numpy_array(box), dtype=np.float64) + bounds = k_grid_bounds_of_cells(cells, k_max) + rows = [] + for f, cell in enumerate(cells): + n_grid = enumerate_k_grid_bounds(bounds[f]).astype(np.float64) + k = n_grid @ (2.0 * math.pi * np.linalg.inv(cell)).T + rows.append(k[np.sum(k * k, axis=-1) <= k_max * k_max]) + vec, mask = _k_blocks(rows) + return _k_grid(box, vec, mask, _frame_slots(n_node), coord) + + +def build_fixed_k_grid( + box: Array, bounds: Array, k_max: float, coord: Array, n_node: Array +) -> KGrid: + r"""Evaluate a fixed integer k grid on the cells of the periodic frames. + + Every frame shares the integer indices within ``bounds`` (see + :func:`enumerate_k_grid_bounds`), so the blocks of a row have the static + count of the grid. The wave numbers of a frame are :math:`\mathbf{k} = + 2\pi B^{-1}\mathbf{n}`, and the points beyond the cutoff + :math:`|\mathbf{k}| \le k_{max}` and the guard points that complete the + last block are masked out. + + A frame whose cell needs indices beyond the grid, a larger cell than the + grid was sized for or a kernel whose cutoff has grown since, is evaluated + on the grid alone: its sum is truncated to the points of the grid, which + is issued as a warning. Its forces and virial remain the derivatives of + that energy. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + bounds + The per-axis bounds :math:`n_d^{max}` of the shared integer grid with + shape ``(3,)``. + k_max + The wave-number cutoff, in Å⁻¹. + coord + Node positions with shape ``(N, 3)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + + Returns + ------- + KGrid + The wave numbers of every frame over the shared grid, one row each. + """ + box = _cells(box) + cells = np.asarray(to_numpy_array(box), dtype=np.float64) + n_grid = enumerate_k_grid_bounds(bounds) + bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(3) + required = k_grid_bounds_of_cells(cells, k_max) + beyond = np.any(required > bounds[None, :], axis=1) + if np.any(beyond): + f = int(np.argmax(beyond)) + warnings.warn( + f"frame {f} needs integer k-grid bounds {required[f].tolist()}, " + f"beyond the fixed grid {bounds.tolist()}: its k-space sum is " + "truncated to the grid; enlarge `k_grid_bounds` for the full sum", + stacklevel=2, + ) + rec = 2.0 * math.pi * np.linalg.inv(cells) + k = np.einsum("jm,fdm->fjd", n_grid.astype(np.float64), rec) + vec, real = _k_blocks(list(k)) + mask = real & (np.sum(vec * vec, axis=-1) <= k_max * k_max) + return _k_grid(box, vec, mask, _frame_slots(n_node), coord) + + +def empty_k_grid(nf: int, coord: Array) -> KGrid: + """The k grid of a batch without periodic frames: guard points and no nodes.""" + xp = array_api_compat.array_namespace(coord) + device = array_api_compat.device(coord) + eye = xp.tile(xp.eye(3, dtype=coord.dtype, device=device)[None, ...], (nf, 1, 1)) + vec, mask = _k_blocks([np.zeros((0, 3), dtype=np.float64)] * nf) + return _k_grid( + eye, vec, mask, np.full((nf, _MIN_ENTRIES), -1, dtype=np.int64), coord + ) + + +def _even_fft_size(n: int) -> int: + """The smallest even integer of at least ``n`` with no prime factor above 7.""" + m = max(2, n + n % 2) + while True: + rest = m + for p in (2, 3, 5, 7): + while rest % p == 0: + rest //= p + if rest == 1: + return m + m += 2 + + +def mesh_dims_of_cells(box: Array, spacing: float) -> tuple[int, int, int]: + r"""The mesh that resolves every cell at ``spacing``. + + Along each lattice vector :math:`\mathbf{a}_d` the mesh holds the + smallest even FFT-friendly number of points (no prime factor above 7) of + at least :math:`\lceil \max_f |\mathbf{a}_d| / h \rceil`, the maximum + running over the frames, so that no frame is evaluated with a spacing + coarser than :math:`h`. The C++ interface sizes the mesh of a frozen + model by the same rule. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + spacing + The largest mesh spacing :math:`h`, in Å. + + Returns + ------- + tuple[int, int, int] + The mesh dimensions. + """ + cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + # The length is formed term by term, as the C++ interface forms it, so + # that both round the ceiling identically. + length = np.sqrt( + cell[..., 0] * cell[..., 0] + + cell[..., 1] * cell[..., 1] + + cell[..., 2] * cell[..., 2] + ) + longest = np.max(length, axis=0) + n1, n2, n3 = (_even_fft_size(math.ceil(float(a) / spacing)) for a in longest) + return n1, n2, n3 + + +def _mesh_modes(xp: object, n: int, dtype: object, device: object) -> Array: + """The integer mode indices of an axis of ``n`` points, in FFT order.""" + m = xp.arange(n, dtype=dtype, device=device) + return xp.where(m > n // 2, m - n, m) + + +def mesh_grid_of_dims( + box: Array, coord: Array, dims: tuple[int, int, int], mask: Array +) -> MeshGrid: + """Pack cells, node positions and a mesh of the given dimensions. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice + vectors. + coord + Node positions with shape ``(N, 3)``. + dims + The even mesh dimensions. + mask + Periodic (1) vs non-periodic (0) frames with shape ``(nf,)``. + + Returns + ------- + MeshGrid + The cells, the mesh and the geometry of the frames. + """ + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + mode_1, mode_2, mode_3 = (_mesh_modes(xp, n, box.dtype, device) for n in dims) + return MeshGrid( + box=box, + volume=xp.astype(_volume(box), box.dtype), + coord=coord, + strain=_zero_strain(box), + mask=mask, + mode_1=mode_1, + mode_2=mode_2, + mode_3=mode_3, + ) + + +def build_mesh_grid(box: Array, spacing: float, coord: Array) -> MeshGrid: + """Pack the cells of the periodic frames and their shared mesh. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + spacing + The largest mesh spacing, in Å (see :func:`mesh_dims_of_cells`). + coord + Node positions with shape ``(N, 3)``. + + Returns + ------- + MeshGrid + The cells, the mesh and the geometry of the frames, every frame + periodic. + """ + xp = array_api_compat.array_namespace(box) + box = _cells(box) + mask = xp.ones((box.shape[0],), dtype=xp.bool, device=array_api_compat.device(box)) + return mesh_grid_of_dims(box, coord, mesh_dims_of_cells(box, spacing), mask) + + +def empty_mesh_grid(nf: int, coord: Array) -> MeshGrid: + """The mesh grid of a batch without periodic frames: every frame masked.""" + xp = array_api_compat.array_namespace(coord) + device = array_api_compat.device(coord) + box = xp.tile(xp.eye(3, dtype=coord.dtype, device=device)[None, ...], (nf, 1, 1)) + mask = xp.zeros((nf,), dtype=xp.bool, device=device) + return mesh_grid_of_dims(box, coord, (2, 2, 2), mask) diff --git a/deepmd/dpmodel/utils/neighbor_graph/topology.py b/deepmd/dpmodel/utils/neighbor_graph/topology.py new file mode 100644 index 0000000000..9a504fc7f2 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/topology.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Topologies an atomic model declares beside its neighbor graph. + +An interaction that the short-range neighbor graph cannot describe reads a +topology of its own, declared by name (see +:meth:`~deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel.required_topologies`): + +- :data:`~.pair_graph.ALL_PAIRS`, the pairs of the non-periodic frames; +- :data:`~.reciprocal_graph.K_GRID` and :data:`~.reciprocal_graph.PME_MESH`, + the reciprocal space of the periodic frames. + +Every topology is a dataclass of arrays that declares which of its fields +carry geometry (``GEOMETRY``), which an autograd backend makes leaves of the +energy (``LEAVES``), and the dynamic axes of every field (``AXES``), so that +casting, re-leafing and the inputs of a compiled or exported graph lower are +derived from the declaration alone. +""" + +from __future__ import ( + annotations, +) + +import dataclasses +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat + +from .pair_graph import ( + ALL_PAIRS, + PairGraph, + build_pair_graph, + check_complete_frames, + empty_pair_graph, +) +from .reciprocal_graph import ( + K_GRID, + PME_MESH, + KGrid, + MeshGrid, + build_fixed_k_grid, + build_k_grid, + build_mesh_grid, + check_periodic_cells, + empty_k_grid, + empty_mesh_grid, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterable, + Sequence, + ) + + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + Topology = PairGraph | KGrid | MeshGrid + +TOPOLOGY_TYPES: dict[str, type] = { + ALL_PAIRS: PairGraph, + K_GRID: KGrid, + PME_MESH: MeshGrid, +} +"""The dataclass of every topology name.""" + + +def _build_reciprocal( + name: str, + box: Array, + coord: Array, + n_node: Array, + params: dict[str, Any], +) -> KGrid | MeshGrid: + """Build a reciprocal-space topology of periodic frames from its parameters.""" + required = {K_GRID: "k_max", PME_MESH: "spacing"}[name] + if required not in params: + raise ValueError( + f"the {name!r} topology requires its `{required}` build parameter" + ) + check_periodic_cells(box) + if name == PME_MESH: + return build_mesh_grid(box, params["spacing"], coord) + if "bounds" in params: + return build_fixed_k_grid(box, params["bounds"], params["k_max"], coord, n_node) + return build_k_grid(box, params["k_max"], coord, n_node) + + +def build_topologies( + names: Iterable[str], + coord: Array, + n_node: Array, + box: Array | None = None, + node_mask: Array | None = None, + *, + atype: Array, + pair_excl: PairExcludeMask | None = None, + topology_params: dict[str, dict[str, Any]] | None = None, +) -> dict[str, Topology]: + """Build the topologies an atomic model declares, beside its neighbor graph. + + A batch is periodic when it carries cells. A declaration of the pairs + alone serves non-periodic batches only. Beside a reciprocal-space + topology, the pairs of a periodic batch come back guard-only and the + reciprocal-space topology of a non-periodic batch masked, so that one + model serves both. + + Parameters + ---------- + names + The declared topology names. + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + n_node + Per-frame node counts with shape ``(nf,)``. + box + Simulation cells with shape ``(nf, 9)`` or ``(nf, 3, 3)``, or + ``None`` for non-periodic systems. + node_mask + Validity of the nodes with shape ``(N,)``, or ``None`` when every + node is valid. + atype + Node types with shape ``(N,)``; invalid nodes carry any type. + pair_excl + The model-level pair exclusion, or ``None``. + topology_params + Build parameters by name: ``{"k_max": ...}`` for the k grid, with + ``"bounds"`` for a fixed integer grid (see + :func:`~.reciprocal_graph.build_fixed_k_grid`), and + ``{"spacing": ...}`` for the mesh. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + Each declared topology under its name; empty when none is declared. + + Raises + ------ + NotImplementedError + If cells are given for a declaration of the pairs alone, or a pair + exclusion for a periodic batch, whose reciprocal-space sum is global. + ValueError + If a name is unknown, or a build parameter is missing. + """ + names = tuple(names) + unknown = [name for name in names if name not in TOPOLOGY_TYPES] + if unknown: + raise ValueError( + f"unknown topology {unknown[0]!r}; expected one of {tuple(TOPOLOGY_TYPES)}" + ) + periodic = any(name in (K_GRID, PME_MESH) for name in names) + if box is None: + nf = n_node.shape[0] + empty = {K_GRID: empty_k_grid, PME_MESH: empty_mesh_grid} + return { + name: ( + build_pair_graph(coord, n_node, atype, node_mask, pair_excl) + if name == ALL_PAIRS + else empty[name](nf, coord) + ) + for name in names + } + if not periodic: + if names: + check_complete_frames(box) + return {} + if pair_excl is not None: + raise NotImplementedError( + "the reciprocal-space sum of a periodic system is global; " + "`pair_exclude_types` cannot exclude pairs from it" + ) + params = topology_params or {} + return { + name: ( + empty_pair_graph(coord) + if name == ALL_PAIRS + else _build_reciprocal(name, box, coord, n_node, params.get(name, {})) + ) + for name in names + } + + +def topology_flat_fields(topology: Topology) -> tuple[Array, ...]: + """The fields of a topology in their declaration order.""" + return tuple( + getattr(topology, field.name) for field in dataclasses.fields(topology) + ) + + +def flatten_topologies(topologies: dict[str, Topology]) -> tuple[Array, ...]: + """Lay topologies out as positional inputs of a graph lower. + + A compiled or exported graph lower receives the fields of every topology, + in the order of the dict and of the declaration of each dataclass, after + its other inputs; :func:`unflatten_topologies` restores them. + + Parameters + ---------- + topologies + The topologies by name. + + Returns + ------- + tuple + The fields of every topology, concatenated. + """ + return tuple( + value + for topology in topologies.values() + for value in topology_flat_fields(topology) + ) + + +def unflatten_topologies( + names: Sequence[str], values: Sequence[Array] +) -> dict[str, Topology]: + """Restore topologies laid out by :func:`flatten_topologies`. + + Parameters + ---------- + names + The topology names, in the order they were laid out. + values + The fields of every topology, concatenated. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + The topologies by name. + + Raises + ------ + ValueError + If a name is unknown, or the number of values does not match the + layout of the names. + """ + topologies = {} + offset = 0 + for name in names: + if name not in TOPOLOGY_TYPES: + raise ValueError(f"unknown topology {name!r}") + cls = TOPOLOGY_TYPES[name] + arity = len(dataclasses.fields(cls)) + if offset + arity > len(values): + break + topologies[name] = cls(*values[offset : offset + arity]) + offset += arity + if len(topologies) != len(names) or offset != len(values): + raise ValueError( + f"{len(values)} flattened values given, which do not lay out the " + f"topologies {tuple(names)}" + ) + return topologies + + +def cast_topology_geometry(topology: Topology, dtype: Any) -> Topology: + """A copy of a topology with its geometry fields cast to ``dtype``.""" + xp = array_api_compat.array_namespace(getattr(topology, topology.GEOMETRY[0])) + return dataclasses.replace( + topology, + **{ + name: xp.astype(getattr(topology, name), dtype) + for name in topology.GEOMETRY + }, + ) diff --git a/deepmd/jax/atomic_model/dp_atomic_model.py b/deepmd/jax/atomic_model/dp_atomic_model.py index 8db31d6d5d..2ef1dc5eb9 100644 --- a/deepmd/jax/atomic_model/dp_atomic_model.py +++ b/deepmd/jax/atomic_model/dp_atomic_model.py @@ -40,6 +40,8 @@ class jax_atomic_model(dpmodel_atomic_model): """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls = None + """No long-range module is implemented in this backend.""" def forward_common_atomic( self, diff --git a/deepmd/pd/model/model/__init__.py b/deepmd/pd/model/model/__init__.py index 59dce4f89b..94897bcfd1 100644 --- a/deepmd/pd/model/model/__init__.py +++ b/deepmd/pd/model/model/__init__.py @@ -64,6 +64,11 @@ def _get_standard_model_components( def get_standard_model(model_params: dict) -> BaseModel: + if model_params.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index d2e0a09e96..b65604c40c 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -208,6 +208,11 @@ def get_linear_model(model_params: dict) -> BaseModel: "sub-model: add an `inner_potential` sub-model to the " "composition instead." ) + if sub.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) if any( sub.get("type") == "inner_potential" for sub in model_params.get("models", []) ): @@ -414,6 +419,11 @@ def get_standard_model(model_params: dict) -> BaseModel: "expands the flag), or spell the composition explicitly with " '`type: "linear_ener"` and an `inner_potential` sub-model.' ) + if model_params.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) diff --git a/deepmd/pt/utils/finetune.py b/deepmd/pt/utils/finetune.py index ad918701da..8a46ae0d69 100644 --- a/deepmd/pt/utils/finetune.py +++ b/deepmd/pt/utils/finetune.py @@ -41,8 +41,8 @@ def _learned_part_path(model: torch.nn.Module) -> str | None: def _move_learned_key(key: str, path: str, source_path: str) -> str: - """Move a descriptor or fitting-net key between learned atomic-model paths.""" - for part in ("descriptor.", "fitting_net."): + """Move a descriptor, fitting-net or long-range key between learned atomic-model paths.""" + for part in ("descriptor.", "fitting_net.", "long_range."): if f".{path}{part}" in key: return key.replace(f".{path}{part}", f".{source_path}{part}", 1) return key diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py index f89b24548d..bc269c6a89 100644 --- a/deepmd/pt_expt/descriptor/dpa4c.py +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -158,6 +158,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.use_amp_infer = use_amp_infer() self._apply_autocast_policy() + def _load_from_state_dict( + self, + state_dict: dict[str, Any], + prefix: str, + *args: Any, + **kwargs: Any, + ) -> None: + # Back-compat: checkpoints written by DeePMD-kit 3.2 hold the gather + # tables of the angular moments, which are formed by broadcasting; + # the stale entries are dropped so that the strict load succeeds. + for name in ("angular_channel_index", "angular_harmonic_index"): + state_dict.pop(prefix + name, None) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + def call_graph( self, graph: Any, diff --git a/deepmd/pt_expt/fitting/__init__.py b/deepmd/pt_expt/fitting/__init__.py index 63bcc6c2ce..2f42a57d48 100644 --- a/deepmd/pt_expt/fitting/__init__.py +++ b/deepmd/pt_expt/fitting/__init__.py @@ -17,6 +17,9 @@ from .invar_fitting import ( InvarFitting, ) +from .latent_charge import ( + LatentChargeFitting, +) from .polarizability_fitting import ( PolarFitting, ) @@ -33,6 +36,7 @@ "DipoleFitting", "EnergyFittingNet", "InvarFitting", + "LatentChargeFitting", "PolarFitting", "PropertyFittingNet", "SeZMEnergyFittingNet", diff --git a/deepmd/pt_expt/fitting/latent_charge.py b/deepmd/pt_expt/fitting/latent_charge.py new file mode 100644 index 0000000000..f4f38d2f0a --- /dev/null +++ b/deepmd/pt_expt/fitting/latent_charge.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting as LatentChargeFittingDP, +) +from deepmd.pt_expt.common import ( + torch_module, +) + + +@torch_module +class LatentChargeFitting(LatentChargeFittingDP): + """Per-atom latent charges predicted from the invariant descriptor (pt_expt).""" diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 2779adfcb0..7b658196f0 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -22,6 +22,14 @@ OutputVariableCategory, OutputVariableDef, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, + build_topologies, + check_complete_frames, + check_periodic_cells, + flatten_topologies, + unflatten_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -721,6 +729,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: ) use_graph_lower = model_uses_graph_lower(model) + topology_names = model.atomic_model.required_topologies() self._dpmodel = model self._is_spin = ( model_params.get("type") == "spin_ener" or "spin" in model_params @@ -775,6 +784,9 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "is_spin": self._is_spin, "lower_input_kind": "graph" if use_graph_lower else "nlist", } + if topology_names: + self.metadata["topologies"] = list(topology_names) + self.metadata["topology_params"] = model.atomic_model.topology_params() if use_graph_lower: from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, @@ -792,7 +804,8 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: # ``communicate_extended_output``, while the graph runner emits local # public keys consumed directly by ``_eval_model_graph``. # - # Graph: (..., source_row_ptr, fparam, aparam, charge_spin) + # Graph: (..., source_row_ptr, fparam, aparam, charge_spin, + # *topology inputs) # Graph spin: (..., source_row_ptr, spin, fparam, aparam, charge_spin) # Nlist: (ext_coord, ext_atype, nlist, mapping, fparam, aparam) # Nlist spin: (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) @@ -862,6 +875,7 @@ def _eager_runner_graph( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None = None, + *topology_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: model_ret = model.forward_common_lower_graph( atype, @@ -879,6 +893,7 @@ def _eager_runner_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + topologies=unflatten_topologies(topology_names, topology_vals), ) return _translate_energy_keys( model_ret, @@ -2137,9 +2152,12 @@ def _eval_model_graph( ) -> tuple[np.ndarray, ...]: """Evaluate generic graph or compact canonical inputs. - Exported artifacts accept a dynamic edge axis. Input preparation uses - the artifact's geometry dtype and layout; both forwards return local - public outputs without extended-atom communication. + Exported artifacts accept a dynamic edge axis, and the inputs of a + model with a long-range module, the fields of every topology that the + ``topologies`` metadata names, follow the graph inputs. Input + preparation uses the artifact's geometry dtype and layout; both + forwards return local public outputs without extended-atom + communication. """ model_inputs, nframes, natoms = self._prepare_graph_inputs( coords, cells, atom_types, fparam, aparam, charge_spin @@ -2297,8 +2315,87 @@ def _prepare_graph_inputs( aparam_t, self._make_charge_spin_input(nframes, charge_spin), ) + topology_names = self.metadata.get("topologies", ()) + if topology_names: + box_input = cells.reshape(nframes, 9) if cells is not None else None + model_inputs = ( + *model_inputs, + *self._long_range_inputs( + topology_names, + coords.reshape(nframes * natoms, 3), + box_input, + atype_t, + graph.n_node, + canonical, + edge_dtype, + DEVICE, + ), + ) return model_inputs, nframes, natoms + def _long_range_inputs( + self, + topology_names: tuple[str, ...], + coord: np.ndarray, + box_input: np.ndarray | None, + atype: torch.Tensor, + n_node: torch.Tensor, + canonical: bool, + edge_dtype: torch.dtype, + device: "torch.device", + ) -> tuple[torch.Tensor | None, ...]: + """Inputs that the lower of a model with a long-range module reads after the graph. + + The compact lower reads the positions and the cells of the frames, and + the zero communicator of frames held whole; a module that serves no + periodic system takes no cell, and the zero cell that marks a + non-periodic frame is never one a caller passes. The graph lower reads + the fields of the topologies, built with the ``topology_params`` of + the archive, or those of the live model of a checkpoint. + """ + nframes = n_node.shape[0] + if canonical: + if PME_MESH not in topology_names: + check_complete_frames(box_input) + elif box_input is not None: + check_periodic_cells(box_input) + box_t = ( + torch.zeros(nframes, 9, dtype=torch.float64, device="cpu") + if box_input is None + else torch.as_tensor( + np.asarray(box_input, dtype=np.float64).reshape(nframes, 9) + ) + ) + return ( + torch.as_tensor(coord, dtype=torch.float64, device=device), + torch.zeros(1, dtype=torch.int64, device="cpu"), + box_t, + ) + # Padding slots of a mixed-type batch are no atoms and join no pair. + # The cell is torch like the coordinates: the topologies it + # namespaces feed the torch graph lower directly (a numpy cell would + # produce numpy leaves). + box_t = ( + torch.as_tensor( + np.asarray(box_input).reshape(nframes, 3, 3), + dtype=edge_dtype, + device=device, + ) + if box_input is not None + else None + ) + topologies = build_topologies( + topology_names, + torch.as_tensor(coord, dtype=edge_dtype, device=device), + n_node, + box_t, + node_mask=atype >= 0, + atype=atype, + pair_excl=self._model_pair_excl(), + topology_params=self.metadata.get("topology_params"), + ) + return flatten_topologies(topologies) + def _build_eval_graph( self, coord_input: np.ndarray, diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical.py b/deepmd/pt_expt/kernels/dpa4c/canonical.py index 77864b6a4a..8fed0b886e 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical.py @@ -58,6 +58,21 @@ def canonical_model_eligible(model: Any) -> bool: return False if getattr(atomic_model, "atom_excl", None) is not None: return False + # A latent-charge long-range module joins the pipeline through the node + # positions and the cells, which its operator reads in place of a pair + # list or a mesh; that pipeline evaluates no pair potential. + if parts[0].has_long_range(): + from deepmd.pt_expt.long_range.latent_charge import ( + LatentChargeLongRange, + ) + + long_range = parts[0].long_range + if not ( + isinstance(long_range, LatentChargeLongRange) + and long_range.operator_eligible() + and parts[1] is None + ): + return False from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( mega_eligible, ) diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py new file mode 100644 index 0000000000..3bb786ff58 --- /dev/null +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compact canonical deployment of compressed DPA4C with a long-range module.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch + +if TYPE_CHECKING: + from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + ) + + +def dpa4c_canonical_long_range_energy_force( + descriptor: Any, + fitting: Any, + long_range: Any, + graph: CanonicalGraph, + atype: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor | None, + box: torch.Tensor | None, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate compressed DPA4C and its long-range module from a compact graph. + + The descriptor operator runs once, and its output feeds both the fused + energy fitting and the charge network of the long-range module. The + derivative of the energy with respect to the descriptor collects the + fitting term and the response of the charges, so a single descriptor + backward yields the edge gradient of both; the long-range force at fixed + charges is added to the assembled force. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor without a spin condition. + fitting + Eligible pt_expt energy fitting network. + long_range + Long-range module whose ``operator_eligible()`` holds. + graph + Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` and + whose ``destination_row_ptr`` has shape ``(N + 1,)`` in ``int64``. + atype + Flat node atom types with shape ``(N,)``. + coord + Node positions with shape ``(N, 3)`` in Å. + communicator + The address of the MPI_Comm handle of a frame decomposed into + domains, one per rank, as an int64 CPU tensor of one element, or zero + or ``None`` for a single domain. Across domains the owned nodes of + every rank make up the frame. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the CPU, + a zero row marking a complete non-periodic frame, or ``None`` when no + frame is periodic. + ownership + Boolean mask selecting energy-contributing nodes with shape ``(N,)``. + atom_bias + Combined atomic energy bias with shape ``(ntypes,)`` in eV. + do_atomic_virial + Whether to return per-node virials. + + Returns + ------- + energy + Per-frame energy with shape ``(F, 1)`` in eV, fp64. + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, fp64. + force + Per-node force with shape ``(N, 3)`` in eV/Å, fp32. + virial + Per-frame virial with shape ``(F, 3, 3)`` in eV, fp32. + atom_virial + Per-node virial with shape ``(N, 3, 3)`` in eV, or an empty tensor. + """ + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + ensure_registered, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + compressed_operator_arguments, + fitting_energy_and_gradient, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + canonical_edge_force_virial, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + ensure_registered as ensure_force_registered, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + frame_scalar_sum, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + ensure_registered as ensure_fitting_registered, + ) + + ensure_registered() + ensure_fitting_registered() + ensure_force_registered() + operator_args = ( + graph.source, + graph.destination_row_ptr, + atype, + *compressed_operator_arguments(descriptor), + int(descriptor.lmax), + *descriptor._compression_scalars, + ) + node_descriptor, state = torch.ops.deepmd.dpa4c_canonical_compress( + graph.edge_vec, *operator_args + ) + # The charge network reads the descriptor before the fused fitting + # overwrites it with its own cotangent. + ( + long_range_energy, + long_range_gradient, + long_range_force, + long_range_atom_virial, + long_range_virial, + ) = long_range.energy_force_from_coord( + node_descriptor, atype, ownership, graph.n_node, coord, communicator, box + ) + atom_energy, descriptor_gradient = fitting_energy_and_gradient( + fitting, node_descriptor, atype, ownership, atom_bias + ) + edge_gradient, _, _ = torch.ops.deepmd.dpa4c_canonical_compress_backward_inplace( + descriptor_gradient + long_range_gradient, + state, + graph.edge_vec, + *operator_args, + ) + force, atom_virial, virial, _ = canonical_edge_force_virial( + edge_gradient, + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + graph.n_node, + graph.edge_vec.new_empty(0), + atype.shape[0], + do_atomic_virial, + ) + # Nodes outside the ownership carry no charge and hence no long-range + # energy; the halo nodes of a periodic frame are images, which the + # reciprocal-space sum holds itself. + atom_energy = atom_energy + long_range_energy + energy = frame_scalar_sum(atom_energy, graph.n_node) + virial = virial + long_range_virial.reshape(-1, 3, 3) + if do_atomic_virial: + atom_virial = atom_virial + long_range_atom_virial.reshape(-1, 3, 3) + return energy, atom_energy, force + long_range_force, virial, atom_virial diff --git a/deepmd/pt_expt/kernels/long_range.py b/deepmd/pt_expt/kernels/long_range.py new file mode 100644 index 0000000000..89db40ad93 --- /dev/null +++ b/deepmd/pt_expt/kernels/long_range.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the operator of the latent-charge long range. + +The operator ``deepmd::latent_charge_sum`` evaluates, for the nodes of +complete frames, the potential that the charges of every other node of the +frame create at each node and the force between the charges. A non-periodic +frame is summed directly over its pairs without listing them, so that the +memory is linear in the node count; a periodic frame is summed over the images +of its cell by the smooth particle-mesh Ewald method, on a mesh sized by the +cell. The CUDA kernels are in ``source/op/pt/long_range``; the CPU kernel +registered here is the reference they are tested against. + +The operator serves inference. It applies the total-charge constraint of the +module, and returns the atomic energy, its derivative with respect to the +charges the charge network predicts, and the force and virial at fixed +charges; the caller adds the response of the charges by back-propagating the +charge gradient through the charge network. + +Given the MPI communicator of a system decomposed into domains, one per rank, +the CUDA kernel evaluates the charged nodes of its rank against those of +every rank; the CPU reference evaluates a single domain. +""" + +import math +from functools import ( + cache, +) +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + frame_id_from_n_node, + unordered_pair_index, +) +from deepmd.pt_expt.kernels.utils import ( + operator_available, +) +from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + SOGKernel, +) + +__all__ = [ + "MAX_CHANNELS", + "MAX_GAUSSIANS", + "MAX_ORDER", + "kernel_eligible", + "kernel_operator_arguments", + "latent_charge_sum", + "op_available", +] + +#: Charge channels and Gaussians the CUDA kernels are instantiated for. +MAX_CHANNELS = 8 +MAX_GAUSSIANS = 32 +#: The largest order of the B-splines of the mesh. +MAX_ORDER = 12 + +_LES = 0 +_SOG = 1 + + +def op_available() -> bool: + """Whether the backend device carries the operator. + + The CPU kernel is the reference of this module, which is registered here + so that the answer does not depend on earlier calls. + """ + # Importing the module loads the operator library. + import deepmd.pt.cxx_op # noqa: F401 + + ensure_registered() + return operator_available("latent_charge_sum") + + +def kernel_eligible(kernel: Any) -> bool: + """Whether the operator evaluates this interaction kernel.""" + if isinstance(kernel, LESKernel): + return True + return isinstance(kernel, SOGKernel) and kernel.amplitude.numel() <= MAX_GAUSSIANS + + +def kernel_operator_arguments(kernel: Any) -> tuple[int, torch.Tensor]: + """Convert an interaction kernel into the code and coefficients of the operator. + + Parameters + ---------- + kernel + A pt_expt interaction kernel accepted by :func:`kernel_eligible`. + + Returns + ------- + tuple[int, torch.Tensor] + The kernel code and its float64 coefficients: ``alpha`` for the erf + kernel; the amplitudes ``a_k`` and then the rates ``c_k = log2(e) / + (2 sigma_k^2)`` of ``a_k 2^(-c_k r^2)`` for the sum of Gaussians. + + Raises + ------ + ValueError + If the operator does not evaluate the kernel. + """ + if not kernel_eligible(kernel): + raise ValueError( + f"the long-range operator does not evaluate {type(kernel).__name__}" + ) + if isinstance(kernel, LESKernel): + coefficient = torch.exp(kernel.log_alpha.to(torch.float64)).reshape(1) + return _LES, coefficient.detach().contiguous() + rate = (0.5 / math.log(2.0)) * torch.exp(-2.0 * kernel.log_width.to(torch.float64)) + coefficient = torch.cat((kernel.amplitude.to(torch.float64), rate)) + return _SOG, coefficient.detach().contiguous() + + +def _kernel_and_slope( + r_sq: torch.Tensor, coefficient: torch.Tensor, kernel: int +) -> tuple[torch.Tensor, torch.Tensor]: + """K(r) and K'(r) / r at the squared distances ``r_sq``.""" + coefficient = coefficient.to(dtype=r_sq.dtype, device=r_sq.device) + if kernel == _LES: + alpha = coefficient[0] + # Coincident nodes, which carry no weight, are evaluated at a short + # distance, where the closed form stays finite. + r_sq = torch.clamp(r_sq, min=1e-12 / alpha**2) + r = torch.sqrt(r_sq) + value = torch.erf(alpha * r) / r + contact = (2.0 / math.sqrt(math.pi)) * alpha + slope = (contact * torch.exp(-(alpha**2) * r_sq) - value) / r_sq + return value, slope + amplitude, rate = coefficient.reshape(2, -1) + weight = torch.exp2(-r_sq[:, None] * rate) + return weight @ amplitude, weight @ (-2.0 * math.log(2.0) * rate * amplitude) + + +def _charges_of_frames( + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + channel: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The charges of the interactions in float64, shifted in the constrained channel. + + Returns the charges with shape ``(N, C)``, the weights of the charged + nodes and the frame of every node with shape ``(N,)``, and the number of + charged nodes of every frame, at least one, with shape ``(F,)``. + """ + n = charge.shape[0] + frame_count = n_node_per_frame.shape[0] + # Padding nodes belong to no frame and carry no charge. + frame = frame_id_from_n_node(n_node_per_frame, n_total=n) + real = torch.arange(n, dtype=torch.int64, device=charge.device) < torch.sum( + n_node_per_frame + ) + weight = (node_mask & real).to(torch.float64) + q = charge.to(torch.float64) * weight[:, None] + count = weight.new_zeros(frame_count).index_add(0, frame, weight) + count = torch.clamp(count, min=1.0) + if channel >= 0: + total = weight.new_zeros(frame_count).index_add(0, frame, q[:, channel]) + shift = (total - total_charge.to(torch.float64)) / count + q = q.clone() + q[:, channel] = q[:, channel] - shift[frame] * weight + return q, weight, frame, count + + +def _direct_reference( + coord: torch.Tensor, + q: torch.Tensor, + n_node_per_frame: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + prefactor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Potential, force and atomic virial over the explicit pairs of every frame.""" + n = coord.shape[0] + pairs = unordered_pair_index(n_node_per_frame) + src, dst = pairs[0], pairs[1] + d = coord[src] - coord[dst] + value, slope = _kernel_and_slope(torch.sum(d * d, dim=-1), coefficient, kernel) + potential = coord.new_zeros(q.shape) + potential = potential.index_add(0, src, q[dst] * value[:, None]) + potential = prefactor * potential.index_add(0, dst, q[src] * value[:, None]) + strength = torch.sum(q[src] * q[dst], dim=-1) * slope + pair_force = -strength[:, None] * d + force = torch.zeros_like(coord).index_add(0, src, pair_force) + force = prefactor * force.index_add(0, dst, -pair_force) + half = (-0.5 * strength)[:, None, None] * d[:, :, None] * d[:, None, :] + atom_virial = coord.new_zeros(n, 3, 3).index_add(0, src, half) + atom_virial = prefactor * atom_virial.index_add(0, dst, half).reshape(n, 9) + return potential, force, atom_virial + + +def _interaction_kernel(coefficient: torch.Tensor, kernel: int) -> Any: + """The dpmodel interaction kernel of the coefficients of the operator.""" + from deepmd.dpmodel.long_range.kernel import LESKernel as LESKernelDP + from deepmd.dpmodel.long_range.kernel import SOGKernel as SOGKernelDP + + values = coefficient.detach().to("cpu", torch.float64).numpy() + if kernel == _LES: + return LESKernelDP(alpha=float(values[0]), precision="float64") + amplitude, rate = values.reshape(2, -1) + width = (0.5 / (math.log(2.0) * rate)) ** 0.5 + return SOGKernelDP( + amplitude=amplitude.tolist(), width=width.tolist(), precision="float64" + ) + + +def _mesh_reference( + coord: torch.Tensor, + q: torch.Tensor, + weight: torch.Tensor, + n_node_per_frame: torch.Tensor, + box: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Potential, force and atomic virial on the mesh of every periodic frame. + + Every frame is evaluated on its own mesh by the evaluation of the training + path, and differentiated with respect to its charges, its positions and + its strain. The atomic virial takes ``F (x) r`` at every node and shares + the rest of the virial of the frame evenly over its charged nodes. + """ + from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, + ) + from deepmd.dpmodel.long_range.reciprocal import ( + mesh_energy, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_mesh_grid, + ) + + interaction = _interaction_kernel(coefficient, kernel) + potential = torch.zeros_like(q) + force = torch.zeros_like(coord) + atom_virial = coord.new_zeros(coord.shape[0], 3, 3) + begin = 0 + for f, length in enumerate(n_node_per_frame.tolist()): + end = begin + length + frame_q = q[begin:end].detach().requires_grad_(True) + frame_coord = coord[begin:end].detach().requires_grad_(True) + with torch.enable_grad(): + mesh = build_mesh_grid(box[f].reshape(1, 3, 3), spacing, frame_coord) + mesh.strain.requires_grad_(True) + frame_id = torch.zeros(length, dtype=torch.int64, device=coord.device) + energy = mesh_energy(frame_q, mesh, frame_id, interaction, order).sum() + energy = energy * (prefactor / COULOMB_CONSTANT) + d_q, d_coord, d_strain = torch.autograd.grad( + energy, (frame_q, frame_coord, mesh.strain) + ) + frame_force = -d_coord + moment = frame_force[:, :, None] * coord[begin:end, None, :] + frame_weight = weight[begin:end] + count = torch.clamp(torch.sum(frame_weight), min=1.0) + rest = (-d_strain[0] - torch.sum(moment, dim=0)) / count + potential[begin:end] = d_q + force[begin:end] = frame_force + atom_virial[begin:end] = moment + rest * frame_weight[:, None, None] + begin = end + return potential, force, atom_virial.reshape(-1, 9) + + +def _reference( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + box: torch.Tensor | None, + coefficient: torch.Tensor, + communicator: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the operator in float64, over the explicit pairs or the meshes.""" + if int(communicator.reshape(-1)[0]) != 0: + raise NotImplementedError( + "the CPU reference of latent_charge_sum evaluates a single domain; " + "a communicator of several ranks requires the CUDA kernel" + ) + frame_count = n_node_per_frame.shape[0] + q, weight, frame, count = _charges_of_frames( + charge, node_mask, n_node_per_frame, total_charge, channel + ) + has_cell = ( + torch.zeros(frame_count, dtype=torch.bool, device="cpu") + if box is None + else torch.any(box != 0.0, dim=-1) + ) + if bool(torch.any(has_cell)) != bool(torch.all(has_cell)): + raise ValueError( + "latent_charge_sum: every frame of a call is periodic, or none is" + ) + if bool(torch.any(has_cell)): + potential, force, atom_virial = _mesh_reference( + coord, + q, + weight, + n_node_per_frame, + box.to(coord.device), + coefficient, + kernel, + prefactor, + order, + spacing, + ) + else: + potential, force, atom_virial = _direct_reference( + coord, q, n_node_per_frame, coefficient, kernel, prefactor + ) + virial = coord.new_zeros(frame_count, 9).index_add(0, frame, atom_virial) + atom_energy = 0.5 * torch.sum(q * potential, dim=-1, keepdim=True) + gradient = potential * weight[:, None] + if channel >= 0: + mean = coord.new_zeros(frame_count).index_add(0, frame, gradient[:, channel]) + gradient = gradient.clone() + gradient[:, channel] = gradient[:, channel] - (mean / count)[frame] * weight + return ( + atom_energy, + gradient.to(charge.dtype), + force.to(charge.dtype), + atom_virial.to(charge.dtype), + virial.to(charge.dtype), + ) + + +def _fake( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + box: torch.Tensor | None, + coefficient: torch.Tensor, + communicator: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + n = coord.shape[0] + return ( + coord.new_empty(n, 1), + torch.empty_like(charge), + charge.new_empty(n, 3), + charge.new_empty(n, 9), + charge.new_empty(n_node_per_frame.shape[0], 9), + ) + + +@cache +def _register() -> None: + torch.library.register_fake("deepmd::latent_charge_sum")(_fake) + library = torch.library.Library("deepmd", "IMPL") + library.impl("latent_charge_sum", _reference, "CPU") + # The registration lives as long as the library object. + _register.library = library + + +def ensure_registered() -> None: + """Register the fake and CPU implementations when the schema is loaded.""" + if isinstance( + getattr(torch.ops.deepmd, "latent_charge_sum", None), + torch._ops.OpOverloadPacket, + ): + _register() + + +def latent_charge_sum( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + kernel: Any, + channel: int, + prefactor: float, + communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Energy, charge gradient, force and virial of the charges of complete frames. + + Parameters + ---------- + coord + Node positions with shape ``(N, 3)`` in Å, float64. The frames occupy + contiguous spans of the node axis; nodes past ``sum(n_node_per_frame)`` + are padding and join no frame. + charge + Raw node charges of the charge network with shape ``(N, C)``, + float32, with ``C`` at most :data:`MAX_CHANNELS`. + node_mask + Nodes that carry charges with shape ``(N,)``, bool; the others carry + none. + n_node_per_frame + Node count of every frame with shape ``(F,)``. + total_charge + Total charge of every frame with shape ``(F,)``, float64, which the + constrained channel sums to. + kernel + The interaction kernel, accepted by :func:`kernel_eligible`. + channel + The constrained channel, which is shifted uniformly over the charged + nodes of every frame to sum to its total charge, or -1 for none. + prefactor + Factor of every pair term, in eV Å. + communicator + The address of the MPI_Comm handle of a system decomposed into + domains, one per rank, as an int64 CPU tensor of one element, or zero, + the default, for a single domain. Across domains the call takes one + frame, whose charged nodes are those of every rank; the outputs are + those of the charged nodes of this rank, the other nodes carry no + charge, and ``virial`` is the share of this rank. + box + The cells of periodic frames with shape ``(F, 9)``, float64 on the + CPU, the rows of every row-major 3x3 cell being its lattice vectors, + or ``None`` for non-periodic frames; a zero row marks a non-periodic + frame. A periodic frame is evaluated on a mesh of spacing at most + ``spacing``, with B-splines of order ``order``. + order + The order of the B-splines of the mesh, at most :data:`MAX_ORDER`. + spacing + The largest mesh spacing, in Å. + + Returns + ------- + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, float64, half of the + interaction of the charges of the node with all others. + charge_gradient + Derivative of the energy with respect to ``charge`` with shape + ``(N, C)``, in eV per unit charge, through the constraint. + force + Force at fixed charges with shape ``(N, 3)``, in eV/Å. + atom_virial + Per-node virial at fixed charges with shape ``(N, 9)`` in eV: half of + every pair term to each of its nodes in a non-periodic frame; in a + periodic frame ``F (x) r`` with the rest of the virial of the cell + shared evenly over the charged nodes. + virial + Per-frame virial at fixed charges with shape ``(F, 9)`` in eV. + """ + ensure_registered() + code, coefficient = kernel_operator_arguments(kernel) + if communicator is None: + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") + return torch.ops.deepmd.latent_charge_sum( + coord, + charge, + node_mask, + n_node_per_frame, + total_charge, + box, + coefficient.to(coord.device), + communicator, + code, + int(channel), + float(prefactor), + int(order), + float(spacing), + ) diff --git a/deepmd/pt_expt/long_range/__init__.py b/deepmd/pt_expt/long_range/__init__.py new file mode 100644 index 0000000000..9c230575e7 --- /dev/null +++ b/deepmd/pt_expt/long_range/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + LESKernel, + SOGKernel, +) +from .latent_charge import ( + LatentChargeLongRange, +) + +__all__ = [ + "BaseLongRange", + "LESKernel", + "LatentChargeLongRange", + "SOGKernel", +] diff --git a/deepmd/pt_expt/long_range/base_long_range.py b/deepmd/pt_expt/long_range/base_long_range.py new file mode 100644 index 0000000000..66135372f9 --- /dev/null +++ b/deepmd/pt_expt/long_range/base_long_range.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import torch + +from deepmd.dpmodel.long_range import ( + make_base_long_range, +) + +BaseLongRange = make_base_long_range(torch.Tensor) diff --git a/deepmd/pt_expt/long_range/kernel.py b/deepmd/pt_expt/long_range/kernel.py new file mode 100644 index 0000000000..3675e6b988 --- /dev/null +++ b/deepmd/pt_expt/long_range/kernel.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.long_range.kernel import LESKernel as LESKernelDP +from deepmd.dpmodel.long_range.kernel import SOGKernel as SOGKernelDP +from deepmd.pt_expt.common import ( + register_dpmodel_mapping, + to_torch_array, +) + + +def _torch_kernel(kernel_dp: type) -> type: + """Wrap a dpmodel interaction kernel as a module whose parameters are trained. + + The attributes listed in ``PARAMETER_NAMES`` are stored as + ``torch.nn.Parameter`` that require a gradient exactly when the kernel is + trainable, both at construction and when a deserialized value is assigned. + """ + + class TorchKernel(kernel_dp, torch.nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + torch.nn.Module.__init__(self) + kernel_dp.__init__(self, *args, **kwargs) + + def __call__(self, *args: Any, **kwargs: Any) -> torch.Tensor: + return torch.nn.Module.__call__(self, *args, **kwargs) + + def forward(self, r_sq: torch.Tensor) -> torch.Tensor: + return self.call(r_sq) + + def __setattr__(self, name: str, value: Any) -> None: + if name in self.PARAMETER_NAMES and "_parameters" in self.__dict__: + if not isinstance(value, torch.nn.Parameter): + value = torch.nn.Parameter( + to_torch_array(value).detach().clone(), + requires_grad=self.trainable, + ) + if name in self._parameters: + self._parameters[name] = value + return + super().__setattr__(name, value) + + TorchKernel.__name__ = TorchKernel.__qualname__ = kernel_dp.__name__ + register_dpmodel_mapping( + kernel_dp, lambda v, cls=TorchKernel: cls.deserialize(v.serialize()) + ) + return TorchKernel + + +LESKernel = _torch_kernel(LESKernelDP) +SOGKernel = _torch_kernel(SOGKernelDP) diff --git a/deepmd/pt_expt/long_range/latent_charge.py b/deepmd/pt_expt/long_range/latent_charge.py new file mode 100644 index 0000000000..b79dffef4c --- /dev/null +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import torch + +from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, +) +from deepmd.dpmodel.long_range.latent_charge import ( + LatentChargeLongRange as LatentChargeLongRangeDP, +) +from deepmd.pt_expt.common import ( + torch_module, +) +from deepmd.pt_expt.fitting.latent_charge import ( + LatentChargeFitting, +) + +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + LESKernel, + SOGKernel, +) + + +@BaseLongRange.register("latent_charge") +@torch_module +class LatentChargeLongRange(LatentChargeLongRangeDP): + """Energy of interacting latent charges, periodic or not (pt_expt).""" + + # The dpmodel class builds the charge network and the kernel, which the + # assignment converts to these pt_expt modules; importing their classes + # registers the conversions. + charge_net: LatentChargeFitting + kernel: LESKernel | SOGKernel + + def _pair_kernel( + self, r_sq: torch.Tensor, src: torch.Tensor, n: int + ) -> torch.Tensor: + # A graph traced for a symbolic pair count plans the reduction of the + # coefficient gradients over the pairs before it sees a batch. In + # training the coefficients therefore reach the pairs through the node + # axis, which turns that reduction into a scatter over the pairs and a + # sum over the nodes, parallel in the pair count of every batch. A + # concrete pair count is reduced directly. + if not (self.training and isinstance(r_sq.shape[0], torch.SymInt)): + return self.kernel(r_sq) + coefficients = { + name: value.expand(n, *value.shape)[src] + for name, value in self.kernel.coefficients(r_sq).items() + } + return self.kernel.evaluate(r_sq, coefficients) + + def operator_eligible(self) -> bool: + """Whether :meth:`energy_force_from_coord` evaluates this module. + + The operator evaluates the kernels, channel counts and B-spline orders + it is instantiated for. The compact lower that calls it reads no frame + input: the charges depend on the descriptor alone, and a constrained + channel sums to zero or to the total charge of a folded charge state. + """ + from deepmd.pt_expt.kernels.long_range import ( + MAX_CHANNELS, + MAX_ORDER, + kernel_eligible, + op_available, + ) + + return ( + op_available() + and kernel_eligible(self.kernel) + and self.charge_net.dim_out <= MAX_CHANNELS + and self.pme_order <= MAX_ORDER + and self.charge_net.numb_fparam == 0 + and self.charge_net.numb_aparam == 0 + and not self.uses_charge_spin() + ) + + def energy_force_from_coord( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + node_mask: torch.Tensor, + n_node: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the module for inference from the positions of the nodes. + + The operator sums a non-periodic frame directly over its pairs, which + it enumerates itself, so that the memory is linear in the node count, + and a periodic frame on the mesh of its cell. It applies the + constraint and returns the force at fixed charges and the derivative + of the energy with respect to the charges the network predicts; + pulled back through the charge network, that derivative is the + derivative of the energy with respect to the descriptor, through which + the caller adds the response of the charges to the positions. + + Parameters + ---------- + descriptor + Invariant node features with shape ``(N, dim_descrpt)``. + atype + Node types with shape ``(N,)``. + node_mask + Nodes carrying charges with shape ``(N,)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + coord + Node positions with shape ``(N, 3)`` in Å. + communicator + The MPI communicator of a frame decomposed into domains, one per + rank, whose charged nodes are those of every rank, in the form of + :func:`~deepmd.pt_expt.kernels.long_range.latent_charge_sum`; + ``None`` for a single domain. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the + CPU, a zero row marking a non-periodic frame, or ``None`` when no + frame is periodic. A module that does not serve periodic systems + evaluates every frame as non-periodic. + + Returns + ------- + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, float64. + descriptor_gradient + Derivative of the energy with respect to the descriptor through + the charges, with the shape and precision of ``descriptor``. + force + Force at fixed charges with shape ``(N, 3)`` in eV/Å, float32. + atom_virial + Per-node virial at fixed charges with shape ``(N, 9)`` in eV, + float32, each pair term split half to each of its nodes. + virial + Per-frame virial at fixed charges with shape ``(nf, 9)`` in eV, + float32; across domains, the share of the nodes of this rank. + """ + from deepmd.pt_expt.kernels.long_range import ( + latent_charge_sum, + ) + + leaf = descriptor.detach().requires_grad_(True) + with torch.enable_grad(): + charge = self.charge_net.call_graph(leaf, atype)["latent_charge"] + total_charge = coord.new_zeros(n_node.shape[0], dtype=torch.float64) + channel = -1 + if self.charge_constraint is not None: + channel = self.charge_constraint["channel"] + total_charge = self._total_charge(n_node.shape[0], total_charge, None, None) + atom_energy, charge_gradient, force, atom_virial, virial = latent_charge_sum( + coord.to(torch.float64).contiguous(), + charge.detach().to(torch.float32).contiguous(), + node_mask.contiguous(), + n_node, + total_charge.contiguous(), + self.kernel, + channel, + COULOMB_CONSTANT, + communicator, + # A module that does not serve periodic systems sums every frame + # directly. + box if self.periodic else None, + self.pme_order, + self.mesh_spacing() if self.periodic else 1.0, + ) + (descriptor_gradient,) = torch.autograd.grad( + charge, leaf, charge_gradient.to(charge.dtype) + ) + return ( + atom_energy, + descriptor_gradient.to(descriptor.dtype), + force, + atom_virial, + virial, + ) diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index a4a497d453..f149053489 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -1,10 +1,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Autograd assembly: graph energy -> force/virial/atom_virial via grad(E, edge_vec). -torch-only. The pure-array scatter (edge_force_virial) is shared with dpmodel; -this module supplies the single backward pass that produces g_e = dE/d(edge_vec). +torch-only. The pure-array scatters (edge_force_virial, pair_force_virial) are +shared with dpmodel; this module supplies the single backward pass that +produces the gradient of the energy with respect to every geometry leaf: the +edge vectors of the neighbor graph and the geometry leaves of the topologies. """ +import dataclasses +from collections.abc import ( + Sequence, +) + import torch from deepmd.dpmodel import ( @@ -19,9 +26,14 @@ get_deriv_name_mag, ) from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, + MeshGrid, NeighborGraph, + PairGraph, edge_force_virial, frame_id_from_n_node, + node_validity_mask, + pair_force_virial, segment_sum, ) from deepmd.pt.utils import ( @@ -44,6 +56,91 @@ ) +def releaf_topologies( + topologies: dict[str, PairGraph | KGrid | MeshGrid], +) -> dict[str, PairGraph | KGrid | MeshGrid]: + """Make the geometry fields of every topology autograd leaves. + + The leaves are made from the inputs of a traced graph lower, so that the + backward roots at its inputs: a leaf created inside a ``make_fx`` trace + has a NaN gradient under inductor. + """ + return { + name: dataclasses.replace( + topology, + **{ + leaf: getattr(topology, leaf).detach().requires_grad_(True) + for leaf in topology.LEAVES + }, + ) + for name, topology in topologies.items() + } + + +def _cell_force_virial( + g_coord: torch.Tensor, + g_strain: torch.Tensor, + coord: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor | None, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + """Force and virial of a reciprocal-space energy from its geometry leaves. + + The gradient with respect to the node positions is the negative force, + and that with respect to the zero strain of every frame the negative + virial. The atomic virial takes ``F_i (x) r_i`` at every node, the force + index first as for the edge virial, and shares the rest of the virial of + the frame, the part of the cell, evenly over its owned nodes, so that the + nodes of a frame sum to its virial; the halo nodes, images that carry no + charge, take no share. + + Parameters + ---------- + g_coord + (N, 3) gradient of the energy with respect to the positions. + g_strain + (nf, 3, 3) gradient of the energy with respect to the strain. + coord + (N, 3) node positions. + n_node + (nf,) per-frame node counts. + n_local + (nf,) per-frame owned node counts, or ``None`` when every node is + owned. + do_atomic_virial + Whether to return the per-node virial. + + Returns + ------- + force + (N, 3) per-node force. + atom_virial + (N, 3, 3) per-node virial, or ``None``. + virial + (nf, 3, 3) per-frame virial. + """ + force = -g_coord + virial = -g_strain + if not do_atomic_virial: + return force, None, virial + n_total = force.shape[0] + frame_id = frame_id_from_n_node(n_node, n_total) + atom_virial = force[:, :, None] * coord.to(force.dtype)[:, None, :] + remainder = virial - segment_sum(atom_virial, frame_id, n_node.shape[0]) + if n_local is None: + owned = node_validity_mask(n_node, n_total) + count = n_node + else: + owned = node_ownership_mask(n_node, n_local, n_total) + count = n_local + share = owned.to(force.dtype) / torch.clamp( + count.to(force.dtype)[frame_id], min=1.0 + ) + atom_virial = atom_virial + remainder[frame_id] * share[:, None, None] + return force, atom_virial, virial + + def edge_energy_deriv( energy: torch.Tensor, edge_vec: torch.Tensor, @@ -60,6 +157,8 @@ def edge_energy_deriv( do_atomic_virial: bool = False, create_graph: bool = False, force_precision: torch.dtype | None = None, + topologies: Sequence[PairGraph | KGrid | MeshGrid] = (), + n_local: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """Return (force, atom_virial_or_None, virial) from a graph energy. @@ -79,6 +178,15 @@ def edge_energy_deriv( reduction. A retained graph (training / double backward) keeps the fp64 leaf dtype so training numerics are unchanged. + The same backward differentiates the energy with respect to the geometry + leaves of every topology in ``topologies``, and their pullbacks are summed + in: the energy reaches the coordinates through all its geometry leaves, so + its complete force and virial are the sum of their pullbacks. The pair + vectors of a pair topology are pulled back by the shared + pair_force_virial scatter; a reciprocal-space topology carries the node + positions, whose gradient is its force, and a zero strain, whose gradient + is its negative virial (see :func:`_cell_force_virial`). + Parameters ---------- energy @@ -109,6 +217,12 @@ def edge_energy_deriv( whether to materialize the per-atom virial (else ``None`` is returned). create_graph whether the backward retains a graph (training, for second-order grad). + topologies + Topologies on the same node axis whose ``LEAVES`` are further + autograd leaves of ``energy``. + n_local + (nf,) per-frame owned node counts, or ``None`` when every node is + owned; the owned nodes share the virial of a cell. Returns ------- @@ -119,9 +233,12 @@ def edge_energy_deriv( virial (nf, 3, 3) per-frame virial (always computed). """ - (g_e,) = torch.autograd.grad( + leaves = [ + getattr(topology, name) for topology in topologies for name in topology.LEAVES + ] + g_e, *g_leaves = torch.autograd.grad( energy.sum() if energy.dim() else energy, - edge_vec, + (edge_vec, *leaves), create_graph=create_graph, retain_graph=True, ) @@ -197,6 +314,29 @@ def edge_energy_deriv( force, atom_virial, virial = edge_force_virial( g_e, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity ) + offset = 0 + for topology in topologies: + grads = g_leaves[offset : offset + len(topology.LEAVES)] + offset += len(topology.LEAVES) + if isinstance(topology, PairGraph): + (g_p,) = grads + leaf_force, leaf_atom_virial, leaf_virial = pair_force_virial( + g_p, + topology.vec, + topology.index, + topology.mask, + n_node, + node_capacity=force.shape[0], + ) + else: + g_coord, g_strain = grads + leaf_force, leaf_atom_virial, leaf_virial = _cell_force_virial( + g_coord, g_strain, topology.coord, n_node, n_local, do_atomic_virial + ) + force = force + leaf_force.to(force.dtype) + virial = virial + leaf_virial.to(virial.dtype) + if do_atomic_virial: + atom_virial = atom_virial + leaf_atom_virial.to(atom_virial.dtype) return force, (atom_virial if do_atomic_virial else None), virial @@ -211,6 +351,7 @@ def fit_output_to_model_output_graph( n_local: torch.Tensor | None = None, force_precision: torch.dtype | None = None, spin_leaf: torch.Tensor | None = None, + topologies: dict[str, PairGraph | KGrid | MeshGrid] | None = None, ) -> dict[str, torch.Tensor]: """Graph analogue of the dense pt_expt ``fit_output_to_model_output``. @@ -282,6 +423,10 @@ def fit_output_to_model_output_graph( ``grad(E, [edge_vec, spin])`` jointly in one call; here the two grads are separate so :func:`edge_energy_deriv`'s signature stays untouched. ``None`` (default): unchanged, no mag output. + topologies + Topologies read by the atomic model, by name. Their geometry leaves + are autograd leaves of ``fit_ret`` next to ``graph.edge_vec``; force + and virial include their pullbacks (see :func:`edge_energy_deriv`). Returns ------- @@ -393,6 +538,8 @@ def reduce_by_frame(data: torch.Tensor) -> torch.Tensor: do_atomic_virial=(vdef.c_differentiable and do_atomic_virial), create_graph=create_graph, force_precision=force_precision if not create_graph else None, + topologies=() if topologies is None else tuple(topologies.values()), + n_local=n_local, ) # 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 9e22a5830b..61c7bf103a 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -19,6 +19,9 @@ from deepmd.dpmodel.model.make_hessian_model import ( make_hessian_model, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + flatten_topologies, +) from deepmd.dpmodel.utils.neighbor_list import ( NeighborList, ) @@ -32,6 +35,7 @@ from .make_model import ( _translate_energy_keys, + check_owned_frames, fused_atom_bias, make_model, ) @@ -146,6 +150,9 @@ def forward_lower_canonical_graph( *, do_atomic_virial: bool, spin: torch.Tensor | None = None, + coord: torch.Tensor | None = None, + communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Evaluate an eligible compressed canonical deployment graph. @@ -177,6 +184,21 @@ def forward_lower_canonical_graph( spin Per-node magnetic moments with shape ``(N, 3)`` for a native-spin model, or ``None``. Ghost rows carry their owner's moment. + coord + Node positions with shape ``(N, 3)`` in Å, float64, for a model + with a long-range module; ``None`` otherwise. + communicator + For a model with a long-range module, the address of the MPI_Comm + handle of a frame decomposed into domains, one per rank, as an + int64 CPU tensor of one element, whose owned nodes are then those + of one domain; zero or ``None`` for a complete frame. + box + For a model with a long-range module, the cells of the frames + with shape ``(nf, 9)``, float64 on the CPU, the rows of every + row-major 3x3 cell its lattice vectors; a zero row marks a + complete non-periodic frame, and ``None`` non-periodic frames. + The halo nodes of a periodic frame are images of its owned + nodes. Returns ------- @@ -217,7 +239,40 @@ def forward_lower_canonical_graph( descriptor = learned.descriptor fitting = learned.fitting_net atom_bias = fused_atom_bias(self.atomic_model, learned) - if use_dpa4c: + if learned.has_long_range(): + from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + dpa4c_canonical_long_range_energy_force, + ) + + if coord is None: + raise ValueError( + "a model with a long-range module reads the node positions" + ) + check_owned_frames(n_node, n_local, communicator, box) + atype, output_mask = self.atomic_model._prepare_graph_nodes( + n_node, + n_local, + atype, + edge_vec, + ) + force_mag = None + mask = output_mask.to(torch.int32) + energy, atom_energy, force, virial, atom_virial = ( + dpa4c_canonical_long_range_energy_force( + descriptor, + fitting, + learned.long_range, + graph, + atype, + coord, + communicator, + box, + output_mask, + atom_bias, + do_atomic_virial, + ) + ) + elif use_dpa4c: # A program without the per-atom virial declines it with the # two-entry placeholders of the source CSR, so that the fused # operator stores no edge cotangent whatever its caller passes. @@ -296,6 +351,9 @@ def forward_lower_canonical_graph_exportable( destination_row_ptr: torch.Tensor, source_row_ptr: torch.Tensor, source_order: torch.Tensor, + coord: torch.Tensor | None = None, + communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, *, do_atomic_virial: bool, **make_fx_kwargs: Any, @@ -308,6 +366,12 @@ def forward_lower_canonical_graph_exportable( Compact graph node and edge tensors. destination_row_ptr, source_row_ptr, source_order Compact dual-CSR topology. + coord, communicator, box + For a model with a long-range module, sample node positions with + shape ``(N, 3)``, float64, a sample communicator, an int64 CPU + tensor of one element, and sample cells with shape ``(nf, 9)``, + float64 on the CPU, which the traced module then takes as its + ninth to eleventh inputs; ``None`` otherwise. do_atomic_virial Whether the traced output includes per-node virial. **make_fx_kwargs @@ -316,9 +380,59 @@ def forward_lower_canonical_graph_exportable( Returns ------- torch.nn.Module - Traced eight-input compact deployment module. + Traced compact deployment module of eight inputs, or eleven with + the node positions, the communicator and the cells. """ model = self + if not (coord is None) == (communicator is None) == (box is None): + raise ValueError( + "the compact lower of a model with a long-range module takes " + "the node positions, the communicator and the cells together" + ) + + if coord is not None: + + def fn_long_range( + 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, + coord: torch.Tensor, + communicator: torch.Tensor, + box: 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, + coord=coord, + communicator=communicator, + box=box, + ) + + return make_fx(fn_long_range, **make_fx_kwargs)( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + coord, + communicator, + box, + ) def fn( atype: torch.Tensor, @@ -638,6 +752,7 @@ def forward_lower_graph_exportable( do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, destination_sorted: bool = False, + topologies: dict[str, Any] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_common_lower_graph`` into an exportable module with @@ -671,6 +786,11 @@ def forward_lower_graph_exportable( and ``destination_order`` is identity. fparam, aparam, do_atomic_virial, charge_spin As in ``forward_lower``. + topologies + The sample topologies the atomic model declares, or ``None``; + their fields follow ``charge_spin`` as further inputs, in the + layout of + :func:`~deepmd.dpmodel.utils.neighbor_graph.flatten_topologies`. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -705,6 +825,7 @@ def forward_lower_graph_exportable( charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, destination_sorted=destination_sorted, + topologies=topologies, **make_fx_kwargs, ) @@ -713,6 +834,67 @@ def forward_lower_graph_exportable( do_grad_r = self.do_grad_r("energy") do_grad_c = self.do_grad_c("energy") + if topologies: + # The topology inputs follow ``charge_spin``, their fields in the + # layout of flatten_topologies. + + def fn_topologies( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + *topology_vals: torch.Tensor, + ) -> dict[str, torch.Tensor]: + model_ret = traced( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + *topology_vals, + ) + return _translate_energy_keys( + model_ret, + do_grad_r=do_grad_r, + do_grad_c=do_grad_c, + do_atomic_virial=do_atomic_virial, + local=True, + ) + + return make_fx(fn_topologies, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + *flatten_topologies(topologies), + ) + def fn( atype: torch.Tensor, n_node: torch.Tensor, diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 2071875b5c..5bd6e68717 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -31,6 +31,9 @@ from deepmd.pt_expt.fitting import ( BaseFitting, ) +from deepmd.pt_expt.long_range import ( + BaseLongRange, +) from deepmd.pt_expt.model.dp_zbl_model import ( DPZBLModel, ) @@ -75,6 +78,7 @@ pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, linear_atomic_model=LinearEnergyAtomicModel, + long_range_base=BaseLongRange, ) get_zbl_model = _model_factory.get_zbl_model @@ -168,6 +172,7 @@ def get_sezm_model(data: dict) -> BaseModel: data["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types) descriptor, fitting, _ = _model_factory.get_model_components(data) + long_range = _model_factory.get_long_range(data, descriptor, fitting) return DPA4EnergyModel( descriptor=descriptor, fitting=fitting, @@ -175,6 +180,7 @@ def get_sezm_model(data: dict) -> BaseModel: atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=pair_exclude_types, preset_out_bias=data.get("preset_out_bias"), + **({} if long_range is None else {"long_range": long_range}), ) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 0dfd77ccc2..25e88487a7 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -23,6 +23,8 @@ from deepmd.dpmodel.utils.neighbor_graph import ( compact_nodes, expand_node_values, + flatten_topologies, + unflatten_topologies, ) from deepmd.pt_expt.common import ( auto_wrapped_class, @@ -43,6 +45,7 @@ from .edge_transform_output import ( fit_output_to_model_output_graph, + releaf_topologies, ) from .transform_output import ( fit_output_to_model_output, @@ -137,7 +140,8 @@ def _fused_energy_force_graph( """ am = model.atomic_model parts = am.fused_decomposition() - if parts is None: + if parts is None or parts[0].has_long_range(): + # The descriptor pipelines evaluate no long-range module. return None learned, pair_potential = parts desc, fit = learned.descriptor, learned.fitting_net @@ -191,6 +195,76 @@ def _fused_energy_force_graph( return ret +def check_owned_frames( + n_node: torch.Tensor, + n_local: torch.Tensor | None, + communicator: torch.Tensor | None = None, + periodic_marks: torch.Tensor | None = None, +) -> None: + """Reject halo nodes in the non-periodic frames of a single domain with a long-range module. + + The module sums over every pair of atoms of a non-periodic frame, so a + frame that one domain holds whole owns all its nodes; halo nodes would + carry no charge and remove their interactions from the sum. Across the + domains of a communicator, every domain owns its share of the frame + instead. The halo nodes of a periodic frame are images of its atoms, + which carry no charge: the reciprocal-space sum holds the images itself. + The counts are read eagerly only: a traced lower receives them from + callers that check them. + + Parameters + ---------- + n_node + Per-frame node counts with shape ``(nf,)``. + n_local + Per-frame owned node counts with shape ``(nf,)``, or ``None``. + communicator + The communicator input of the compact lower, a CPU tensor holding the + address of an MPI_Comm handle or zero, or ``None`` for a single domain. + periodic_marks + A tensor with a nonzero entry when the frames are periodic, such as + their cells or the frame mask of a reciprocal-space topology, or + ``None`` for non-periodic frames. It is read only after the fake + check, so that a traced lower computes nothing on it. + + Raises + ------ + ValueError + If a non-periodic frame of a single domain carries halo nodes. + """ + from torch._subclasses.fake_tensor import ( + is_fake, + ) + + if n_local is None or n_local is n_node or is_fake(n_local): + return + if communicator is not None and int(communicator.reshape(-1)[0]) != 0: + return + if periodic_marks is not None and bool(torch.any(periodic_marks != 0)): + return + if not torch.equal(n_local, n_node): + raise ValueError( + "a model with a long-range module evaluates complete non-periodic " + "frames on a single domain, whose graph owns every node (n_local " + f"== n_node), but got n_local {n_local.tolist()} for n_node " + f"{n_node.tolist()}" + ) + + +def _reciprocal_masks(topologies: dict[str, Any]) -> torch.Tensor | None: + """The mask of the reciprocal-space topology, true on periodic entries, or ``None``.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + ) + + masks = [ + topology.mask + for topology in topologies.values() + if not isinstance(topology, PairGraph) + ] + return masks[0] if masks else None + + def _pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor: """Append a single ``-1`` column to ``nlist`` for export-time tracing. @@ -356,7 +430,8 @@ class _WrapperForwardEnergyGraph: """Graph twin of :class:`_WrapperForwardEnergy` for the Hessian. Given flattened LOCAL coordinates for one frame, rebuilds the carry-all - ``NeighborGraph`` (so ``edge_vec`` tracks the coordinates through + ``NeighborGraph`` and the topologies the atomic model declares (so + ``edge_vec`` and the pair vectors track the coordinates through ``autograd.functional.hessian``'s double-backward) and returns the scalar reduced-output component. Unlike the dense wrapper this differentiates w.r.t. the ``(nloc, 3)`` LOCAL coordinates directly -- the carry-all graph @@ -399,6 +474,15 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: ng = build_neighbor_graph_for_method( self.method, cc, self.atype, self.box, self.rcut, self.pair_excl ) + # The topologies follow the same coordinates, so the Hessian holds + # every geometry path of the energy and the blocks coupling them. + topologies = self.model.atomic_model.build_topologies( + coord_flat.reshape(self.nloc, 3), + ng.n_node, + self.box, + atype=self.atype.reshape(-1), + pair_excl=self.pair_excl, + ) atomic_ret = self.model.atomic_model.forward_common_atomic_graph( ng, self.atype.reshape(-1), @@ -406,6 +490,7 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: aparam=self.aparam, spin=self.spin, charge_spin=self.charge_spin, + topologies=topologies, ) # atomic_ret[kk]: flat (N, *def), N == nloc for a single-frame carry-all # graph (all nodes owned); reduced output = sum over the node axis. @@ -608,6 +693,7 @@ def forward_common_lower_graph( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, torch.Tensor]: """Graph-native lower with autograd force/virial (dpa1/se_atten concat-tebd, attention included). @@ -696,6 +782,13 @@ def forward_common_lower_graph( DescrptBlockRepformers._exchange_ghosts_graph`). Only meaningful for descriptors whose ``has_message_passing_across_ranks()`` is ``True``. + topologies + The topologies the atomic model declares by + ``required_topologies()``, built on the node axis of the + graph. The geometry fields every topology names in its + ``LEAVES`` become autograd leaves next to ``edge_vec``, and + force and virial include their pullbacks. ``None`` (default) + for models declaring none. Returns ------- @@ -712,6 +805,17 @@ def forward_common_lower_graph( # make edge_vec the autograd leaf for the energy backward edge_vec = edge_vec.detach().requires_grad_(True) + if topologies: + if comm_dict is not None: + raise NotImplementedError( + "the long-range topologies describe complete systems; " + "a domain-decomposed system is not supported" + ) + check_owned_frames( + n_node, n_local, periodic_marks=_reciprocal_masks(topologies) + ) + # Every topology carries geometry leaves of its own. + topologies = releaf_topologies(topologies) if spin is not None: # second autograd leaf: force_mag = -dE/dspin (native spin). # Deliberately a SEPARATE leaf/backward from edge_vec rather @@ -751,6 +855,7 @@ def forward_common_lower_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, ) # ``forward_common_atomic_graph`` returns flat ``(N, *)`` output. # Pass directly to the flat-N transform; no rectangular reshape needed. @@ -772,6 +877,7 @@ def forward_common_lower_graph( # dynamic-edge torch.export. See fit_output_to_model_output_graph. node_capacity=atype.shape[0], n_local=n_local, + topologies=topologies, ) def _resolve_graph_method( @@ -892,6 +998,7 @@ def call_common_ragged( method = getattr(self, "neighbor_graph_method", None) if method is None: method = resolve_neighbor_graph_method("auto", coord.device) + pair_excl = getattr(self.atomic_model, "pair_excl", None) graph = build_ragged_neighbor_graph( method, coord, @@ -899,7 +1006,14 @@ def call_common_ragged( n_node, box, self.get_rcut(), - getattr(self.atomic_model, "pair_excl", None), + pair_excl, + ) + topologies = self.atomic_model.build_topologies( + coord.detach(), + graph.n_node, + box, + atype=atype, + pair_excl=pair_excl, ) predict = self.forward_common_lower_graph( atype, @@ -918,6 +1032,7 @@ def call_common_ragged( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) # The per-atom mask a rectangular batch carries is what tells the # loss each frame's real atom count. Nothing here is padded, so the @@ -1022,6 +1137,13 @@ def _call_common_graph( spin_flat = ( spin.reshape(n_padded, 3)[node_index] if spin is not None else None ) + topologies = self.atomic_model.build_topologies( + cc.detach().reshape(n_padded, 3)[node_index], + ng.n_node, + bb, + atype=atype_flat, + pair_excl=pair_excl, + ) model_predict = self.forward_common_lower_graph( atype_flat, ng.n_node, @@ -1039,6 +1161,7 @@ def _call_common_graph( aparam=ap_flat, spin=spin_flat, charge_spin=charge_spin, + topologies=topologies, ) # ``forward_common_lower_graph`` returns flat ``(N, *)`` per-atom # outputs over the real atoms. Scatter them back onto the padded @@ -1249,6 +1372,7 @@ def forward_common_lower_graph_exportable( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, destination_sorted: bool = False, + topologies: dict[str, Any] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """make_fx trace of ``forward_common_lower_graph`` with ``edge_vec`` @@ -1289,6 +1413,13 @@ def forward_common_lower_graph_exportable( output. ``None`` (default) is the existing, unconditioned trace with no mag output -- used by every non-spin caller of this generic (descriptor-agnostic) exportable. + topologies + Sample topologies the atomic model declares by + ``required_topologies()``, or ``None``. When given, they are + traced inputs laid out by + :func:`~deepmd.dpmodel.utils.neighbor_graph.flatten_topologies` + after ``charge_spin``; a model reading topologies carries no + native spin. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -1299,8 +1430,9 @@ def forward_common_lower_graph_exportable( A traced module whose ``forward`` accepts ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, fparam, aparam, charge_spin, spin)`` and - returns a dict with the same internal keys as + source_row_ptr, fparam, aparam, charge_spin, spin)`` -- or, + with topologies, ``spin`` replaced by the topology inputs -- + and returns a dict with the same internal keys as ``forward_common_lower_graph``. """ validate_graph_csr_for_export( @@ -1315,6 +1447,61 @@ def forward_common_lower_graph_exportable( ) model = self + if topologies: + topology_names = tuple(topologies) + + def fn_topologies( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + *topology_vals: torch.Tensor, + ) -> dict[str, torch.Tensor]: + return model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=destination_sorted, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + topologies=unflatten_topologies(topology_names, topology_vals), + ) + + return make_fx(fn_topologies, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + *flatten_topologies(topologies), + ) + # ``spin`` is a traced INPUT only when a real spin tensor is # given (native-spin models). For every non-spin caller # (``spin is None``) the traced ``fn`` must have the SAME arity as diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 051701dc96..94b0134aae 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -81,6 +81,10 @@ from deepmd.dpmodel.utils.learning_rate import ( make_learning_rate_schedule, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + flatten_topologies, + unflatten_topologies, +) from deepmd.dpmodel.utils.training_utils import ( compute_total_numb_batch, ) @@ -129,6 +133,9 @@ from deepmd.pt_expt.model import ( get_model, ) +from deepmd.pt_expt.model.edge_transform_output import ( + releaf_topologies, +) from deepmd.pt_expt.model.graph_lower import ( model_uses_graph_lower, ) @@ -317,6 +324,10 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: atomic model and fitting. Shared weights do not make those branches interchangeable: tasks without a preset bias disable the reference, while a folded reference no longer appends isolated atoms to the graph. + + A long-range module is part of the traced computation and of its input + schema, so its identity (that of its first parameter) closes the key; + tasks share a graph only when they share the module as well. """ descriptor_id: int = 0 try: @@ -342,7 +353,17 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: break except AttributeError: pass - return (int(model.has_spin()), descriptor_id, fitting_id, *vacuum_state) + long_range_id = 0 + long_range = getattr(getattr(model, "atomic_model", None), "long_range", None) + if long_range is not None: + long_range_id = next((id(p) for p in long_range.parameters()), id(long_range)) + return ( + int(model.has_spin()), + descriptor_id, + fitting_id, + *vacuum_state, + long_range_id, + ) # --------------------------------------------------------------------------- @@ -785,6 +806,12 @@ def _trace_and_compile_graph( autograd leaf. Their compiled lower returns ``force_mag`` and ``mask_mag`` through the same model-owned translation used by eager execution. + A model declaring topologies (``atomic_model.required_topologies()``) + receives the fields of each, laid out by :func:`flatten_topologies`, as + further inputs ahead of the task buffers; their dynamic axes, the pairs, + the wave numbers or the axes of the mesh, are symbols distinct from the + other three. + Parameters ---------- model @@ -862,6 +889,7 @@ def _trace_and_compile_graph( # its graph builders omit the CSR metadata carried by deployment inputs. from deepmd.pt_expt.utils.serialization import ( build_synthetic_graph_inputs, + build_synthetic_topologies, check_graph_trace_torch_version, count_synthetic_graph_edges, ) @@ -919,6 +947,22 @@ def _trace_and_compile_graph( # tensors, while forward_common_lower_graph takes it after them. s_spin = s_conditioning.pop(0) if spin is not None else None s_fparam, s_aparam, s_charge_spin = s_conditioning + # The dynamic axes of the topologies take sample lengths distinct from the + # other three and from the CSR offsets. + s_topologies = build_synthetic_topologies( + model, + trace_N, + _next_safe_prime( + e_max + 2, _forbidden | {trace_nf, trace_N, trace_N + 1, e_max} + ), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=_trace_device, + nframes=trace_nf, + forbidden=_forbidden | {e_max}, + ) + topology_names = tuple(s_topologies) + s_topology_vals = flatten_topologies(s_topologies) + n_topology_vals = len(s_topology_vals) def fn( atype: torch.Tensor, @@ -935,8 +979,10 @@ def fn( aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, spin: torch.Tensor | None, - *task_buf_vals: torch.Tensor, + *extra_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: + topologies = unflatten_topologies(topology_names, extra_vals[:n_topology_vals]) + task_buf_vals = extra_vals[n_topology_vals:] # Patch task-specific buffers with the proxy tensors so make_fx records # them as FX placeholders (mirrors the dense ``_trace_and_compile``). originals: dict[str, torch.Tensor | None] = {} @@ -970,6 +1016,7 @@ def fn( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) if spin is None: return _translate_energy_keys( @@ -1016,6 +1063,7 @@ def fn( s_aparam, s_charge_spin, s_spin, + *s_topology_vals, *task_buf_vals_trace, ) @@ -1660,6 +1708,15 @@ def _forward_graph( aparam = aparam[node_index] if spin is not None: spin = spin[node_index] + # The topologies the atomic model declares are built here as well, + # outside the compiled region, on the node axis of the graph. + topologies = _model.atomic_model.build_topologies( + coord_3d if ragged else coord_3d.reshape(n_padded, 3)[node_index], + ng.n_node, + box_flat, + atype=atype_flat, + pair_excl=pair_excl, + ) # Lazy compile of the GRAPH lower (cached per structure key and mode). compiled_lower = self._compiled_lower_by_mode.get(self.training) @@ -1687,6 +1744,9 @@ def _forward_graph( edge_vec = ng.edge_vec.detach().requires_grad_(True) if spin is not None: spin = spin.detach().requires_grad_(True) + # The topologies enter as their fields, the geometry of every one a + # further leaf. + topology_vals = flatten_topologies(releaf_topologies(topologies)) if self._task_buf_order: try: @@ -1728,6 +1788,7 @@ def _forward_graph( aparam, charge_spin, spin, + *topology_vals, *task_buf_vals, ) self._report_pending_compile() @@ -2283,6 +2344,9 @@ def initialize_statistics( pretrained_state = pretrained_wrapper.state_dict() target_state = self._unwrapped.state_dict() new_state = target_state.copy() + # The atomic models, by module path, whose long-range module + # takes its initialization rather than pretrained weights. + fresh_long_range: dict[str, str] = {} for model_key in self.model_keys: finetune_rule = finetune_links[model_key] source_key = finetune_rule.get_model_branch() @@ -2312,11 +2376,6 @@ def initialize_statistics( for key in target_state: if not key.startswith(prefix): continue - if ( - finetune_rule.get_random_fitting() - and ".descriptor." not in key - ): - continue pretrained_key = f"model.{source_key}." + key[len(prefix) :] if ( pretrained_key not in pretrained_state @@ -2326,10 +2385,39 @@ def initialize_statistics( pretrained_key = _move_learned_key( pretrained_key, path, source_path ) - if pretrained_key in pretrained_state: + use_random = ( + finetune_rule.get_random_fitting() + and ".descriptor." not in key + ) + if use_random or pretrained_key not in pretrained_state: + if ".long_range." in key: + owner = key[: key.index(".long_range.")] + fresh_long_range[owner] = model_key + else: new_state[key] = pretrained_state[pretrained_key] self._unwrapped.load_state_dict(new_state) + # A long-range module added to the pretrained model takes the + # input statistics of the fine-tuning data, as in training from + # scratch, before the bias adjustment evaluates the model. + for owner, model_key in fresh_long_range.items(): + atomic_model = self._unwrapped.get_submodule(owner) + + def long_range_stat( + _atomic_model: torch.nn.Module = atomic_model, + _model_key: str = model_key, + ) -> None: + _atomic_model.compute_long_range_input_stat( + self._sample_funcs[_model_key] + ) + + self._run_stat_on_chief( + long_range_stat, + operation=f"long-range statistics for task {model_key!r}", + ) + if self.is_distributed: + self._broadcast_model_stat(atomic_model.long_range) + # Per-branch bias adjustment (rank 0 only, then broadcast) for model_key in self.model_keys: finetune_rule = finetune_links[model_key] diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 7ddd2589f6..f7d9ad2511 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import contextlib import ctypes +import dataclasses import json import logging import os @@ -14,6 +15,9 @@ import numpy as np import torch +from deepmd.dpmodel.utils.neighbor_graph import ( + flatten_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -688,18 +692,120 @@ def build_synthetic_graph_inputs( ) +def build_synthetic_topologies( + model: torch.nn.Module, + n_node_total: int, + n_entry: int, + *, + dtype: torch.dtype, + device: torch.device, + nframes: int, + forbidden: set[int] | frozenset[int] = frozenset(), +) -> dict[str, Any]: + """Build trace-time samples of the topologies a model declares. + + Single source of the trace-time topology inputs, shared by graph export + and compiled training. Tracing reads only their shapes. The pair + topology lists ``n_entry`` pairs, a length the caller keeps distinct from + the other dynamic axes; a mesh takes the three smallest even dimensions, + and a k grid the three smallest lengths for its slots, its blocks and the + points of a block, of at least four that no other axis and no entry of + ``forbidden`` has. Every length thus differs from the two rows of a pair + index and the three components of a vector, and from every other dynamic + axis, which a compiled graph would otherwise identify with it. + + Parameters + ---------- + model : torch.nn.Module + The model; its atomic model declares the topologies. + n_node_total : int + Length of the node axis. + n_entry : int + Length of the pair axis. + dtype : torch.dtype + Precision of the geometry. + device : torch.device + Target device. + nframes : int + Length of the frame axis. + forbidden : set[int] + Sizes of the static dimensions of the model, which the mesh and k-grid + axes avoid. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + A sample of every declared topology; empty when none is declared. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + K_GRID, + PME_MESH, + KGrid, + PairGraph, + ) + from deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph import ( + mesh_grid_of_dims, + ) + + taken = set(forbidden) | {nframes, n_node_total, n_node_total + 1, n_entry} + mesh_dims = [] + for size in range(4, 4 + 2 * len(taken) + 6, 2): + if size not in taken and len(mesh_dims) < 3: + mesh_dims.append(size) + free = sorted(set(range(4, len(taken) + 10)) - taken - set(mesh_dims)) + n_slot, n_block, n_point = free[:3] + coord = torch.ones(n_node_total, 3, dtype=dtype, device=device) + # Contiguous like the inputs at run time: a compiled graph reads the + # strides of its samples. + box = torch.eye(3, dtype=dtype, device=device).repeat(nframes, 1, 1) + topologies = {} + for name in model.atomic_model.required_topologies(): + if name == K_GRID: + slot = torch.arange(nframes * n_slot, dtype=torch.int64, device=device) + blocks = (nframes, n_block, n_point) + topologies[name] = KGrid( + # Nonzero wave numbers: the kernel transform carries 1/k^2. + vec=torch.ones(*blocks, 3, dtype=dtype, device=device), + mask=torch.ones(*blocks, dtype=torch.bool, device=device), + slot=(slot % n_node_total).reshape(nframes, n_slot), + volume=torch.ones(nframes, dtype=dtype, device=device), + coord=coord, + strain=torch.zeros(nframes, 3, 3, dtype=dtype, device=device), + ) + elif name == PME_MESH: + topologies[name] = mesh_grid_of_dims( + box, + coord, + tuple(mesh_dims), + torch.ones(nframes, dtype=torch.bool, device=device), + ) + else: + index = torch.arange(2 * n_entry, dtype=torch.int64, device=device) + topologies[name] = PairGraph( + index=index.reshape(2, n_entry) % n_node_total, + vec=torch.ones(n_entry, 3, dtype=dtype, device=device), + mask=torch.ones(n_entry, dtype=torch.bool, device=device), + ) + return topologies + + def build_synthetic_canonical_graph_inputs( model: torch.nn.Module, e_max: int, *, device: torch.device, want_spin: bool = False, + want_long_range: bool = False, ) -> tuple[torch.Tensor, ...]: """Build compact canonical trace inputs for compressed CUDA descriptors. ``want_spin`` appends the per-node moment at slot 8, the last slot of the compact ABI, matching :meth:`~deepmd.pt_expt.model.native_spin_model.NativeSpinEnergyModel.forward_lower_canonical_graph_exportable`. + ``want_long_range`` appends the inputs that the compact lower of a model + with a long-range module reads: the float64 node positions at slot 8, the + int64 communicator, a CPU tensor of one element, at slot 9, and the + float64 cells of the frames, a CPU tensor of shape ``(nf, 9)``, at slot 10. """ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, @@ -757,11 +863,31 @@ def build_synthetic_canonical_graph_inputs( compact.source_row_ptr, compact.source_order, ) - return compact_sample if spin is None else (*compact_sample, spin) - + if spin is not None: + return (*compact_sample, spin) + if want_long_range: + # Tracing reads the shape of the positions alone; distinct values keep + # the sample nodes apart. The communicator of a single domain is zero. + coord = 0.7 * torch.arange( + 3 * atype.shape[0], dtype=torch.float64, device=device + ).reshape(-1, 3) + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") + box = torch.zeros(n_node.shape[0], 9, dtype=torch.float64, device="cpu") + return (*compact_sample, coord, communicator, box) + return compact_sample + + +def _build_canonical_graph_dynamic_shapes( + *sample_inputs: torch.Tensor, +) -> tuple: + """Build dynamic node, edge and index-word axes for the compact ABI. -def _build_canonical_graph_dynamic_shapes(*, want_spin: bool) -> tuple: - """Build dynamic node, edge and index-word axes for the compact ABI.""" + The trailing slots, the moments of a native-spin model or the positions, + the communicator and the cells of a model with a long-range module, are + present only for such models, so the sample length selects the shape + tuple. The communicator holds one address and is static; the cells follow + the frame axis. + """ 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) @@ -780,7 +906,10 @@ def _build_canonical_graph_dynamic_shapes(*, want_spin: bool) -> tuple: {0: source_row_dim}, {0: source_order_dim}, ) - return (*shapes, {0: node_dim}) if want_spin else shapes + trailing = ({0: node_dim}, None, {0: nframes_dim})[ + : len(sample_inputs) - len(shapes) + ] + return (*shapes, *trailing) def count_synthetic_graph_edges( @@ -833,9 +962,53 @@ def count_synthetic_graph_edges( return int(edge_mask.sum()) +def _topology_dynamic_shapes( + topology_names: tuple[str, ...], shared: dict[str, Any] +) -> list[dict[int, Any]]: + """Dynamic-shape specifications of the fields of the topologies of a graph lower. + + Every field follows the ``AXES`` its dataclass declares. The frame and + node axes are those of the other inputs in ``shared``; the entry axis of + every topology is a dimension of its own, of at least the two guard + entries, and so are the slot axis and the point axis of a k grid; a mesh + axis is twice a dimension, since the mesh dimensions are even. + + Parameters + ---------- + topology_names : tuple[str, ...] + The topologies, in the order of their inputs. + shared : dict[str, Any] + The dimensions of the ``"frame"`` and ``"node"`` axes. + + Returns + ------- + list[dict[int, Any]] + One specification per field, in the layout of flatten_topologies. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + TOPOLOGY_TYPES, + ) + + specs = [] + for name in topology_names: + cls = TOPOLOGY_TYPES[name] + dims = dict(shared) + dims["entry"] = torch.export.Dim(f"{name}_entries", min=2) + dims["slot"] = torch.export.Dim(f"{name}_slots", min=2) + dims["k_point"] = torch.export.Dim(f"{name}_block_points", min=2) + for axis in ("mesh_1", "mesh_2", "mesh_3"): + dims[axis] = 2 * torch.export.Dim(f"{name}_{axis}_half", min=1) + specs += [ + {dim: dims[axis] for dim, axis in cls.AXES[field.name].items()} + for field in dataclasses.fields(cls) + ] + return specs + + def _build_graph_dynamic_shapes( *sample_inputs: torch.Tensor | None, is_native_spin: bool = False, + topology_names: tuple[str, ...] = (), ) -> tuple: """Build dynamic-shape specifications for the graph-form forward_lower export. @@ -863,6 +1036,9 @@ def _build_graph_dynamic_shapes( Whether ``sample_inputs`` follows the native-spin positional ABI (spin at slot 10) instead of the regular energy ABI (charge_spin at slot 12). + topology_names : tuple[str, ...] + Topologies whose fields follow the inputs above (see + :func:`_topology_dynamic_shapes`). """ nframes_dim = torch.export.Dim("nframes", min=1) n_node_total_dim = torch.export.Dim("n_node_total", min=1) @@ -879,6 +1055,9 @@ def _build_graph_dynamic_shapes( {0: nedge_dim}, # source_order: (E,) {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) ) + topologies = _topology_dynamic_shapes( + topology_names, {"frame": nframes_dim, "node": n_node_total_dim} + ) if is_native_spin: spin = sample_inputs[10] fparam = sample_inputs[11] @@ -892,6 +1071,7 @@ def _build_graph_dynamic_shapes( {0: nframes_dim} if fparam is not None else None, # fparam {0: n_node_total_dim} if aparam is not None else None, # aparam {0: nframes_dim} if charge_spin is not None else None, # charge_spin + *topologies, ) fparam = sample_inputs[10] aparam = sample_inputs[11] @@ -904,6 +1084,7 @@ def _build_graph_dynamic_shapes( # dim would make torch.export prove/reject the equality). {0: n_node_total_dim} if aparam is not None else None, # aparam {0: nframes_dim} if charge_spin is not None else None, # charge_spin + *topologies, ) @@ -1293,6 +1474,16 @@ def _probe_has_message_passing(obj: object) -> bool | None: pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] break meta["pair_exclude_types"] = pair_exclude_types + # Topologies the lower reads beside the neighbor graph, and the parameters + # they are built with; absent when the model reads none. The graph lower + # takes the fields of each after its other inputs (see + # ``flatten_topologies``), built by the caller from the coordinates and + # the cell of each frame. The compact canonical lower takes the node + # positions instead, from which its operators evaluate the long range. + topologies = model.atomic_model.required_topologies() + if topologies: + meta["topologies"] = list(topologies) + meta["topology_params"] = model.atomic_model.topology_params() return meta @@ -1961,6 +2152,15 @@ def _trace_and_export_impl( # the deserialized model resides, so it precedes the move below. _fold_case_embeddings(model) model.eval() + # A frozen model evaluates the reciprocal sum of a periodic frame on the + # mesh of its cell, whatever the mode it was trained in. + from deepmd.dpmodel.long_range.latent_charge import ( + LatentChargeLongRange, + ) + + for module in model.modules(): + if isinstance(module, LatentChargeLongRange): + module.use_mesh_evaluation() # Fold the vacuum reference on the CPU tracing device. vacuum_descriptor # evaluates an isolated atom of every type through the descriptor; on CUDA # that path can allocate several GiB even when export already holds most of @@ -2130,11 +2330,17 @@ def _trace_and_export_impl( while e_sample in (_forbidden | {nframes_sample, n_sample}): e_sample += 1 if canonical: + # The compact lower of a model with a long-range module reads the + # node positions and the cells, from which its operator evaluates + # the long range, and the communicator of the domains of a + # decomposed frame. + want_long_range = bool(model.atomic_model.required_topologies()) sample_inputs = build_synthetic_canonical_graph_inputs( model, e_sample, device=torch.device("cpu"), want_spin=is_native_spin, + want_long_range=want_long_range, ) traced = model.forward_lower_canonical_graph_exportable( *sample_inputs, @@ -2142,9 +2348,7 @@ def _trace_and_export_impl( tracing_mode="symbolic", _allow_non_fake_inputs=True, ) - dynamic_shapes = _build_canonical_graph_dynamic_shapes( - want_spin=is_native_spin - ) + dynamic_shapes = _build_canonical_graph_dynamic_shapes(*sample_inputs) elif with_comm_dict: # Load libdeepmd_op_pt.so and register border_op fake/autograd # metadata now, mirroring the dense with-comm precedent below. @@ -2218,6 +2422,20 @@ def _trace_and_export_impl( device=torch.device("cpu"), want_spin=is_native_spin, ) + # The pair axis is dynamic; its sample length is kept distinct + # from the other axes and from the CSR offsets. + topologies = build_synthetic_topologies( + model, + n_sample, + next_safe_prime( + e_sample + 2, + _forbidden | {nframes_sample, n_sample, n_sample + 1, e_sample}, + ), + dtype=edge_dtype, + device=torch.device("cpu"), + nframes=nframes_sample, + forbidden=_forbidden | {e_sample}, + ) if is_native_spin: # Native-spin ABI (NativeSpinEnergyModel.forward_lower_graph_exportable): # slot 10 is ``spin`` (mandatory), slots 11/12 are @@ -2243,12 +2461,16 @@ def _trace_and_export_impl( do_atomic_virial=do_atomic_virial, charge_spin=sample_inputs[12], destination_sorted=True, + topologies=topologies, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) dynamic_shapes = _build_graph_dynamic_shapes( - *sample_inputs, is_native_spin=is_native_spin + *sample_inputs, + is_native_spin=is_native_spin, + topology_names=tuple(topologies), ) + sample_inputs = (*sample_inputs, *flatten_topologies(topologies)) output_keys = traced_output_keys(traced) exported = torch.export.export( traced, @@ -2579,7 +2801,81 @@ def _charge_state_descriptor(data: dict, metadata: dict) -> Any | None: return descriptor -def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: +def _charge_state_long_range(data: dict, metadata: dict) -> Any | None: + """Return the long-range module folded to the charge state, if any. + + Compressing the atomic model folds a long-range charge constraint that + reads its total charge from the charge state to the condition of the + descriptor, so the fold rebuilds the state of the module with the + artifacts of the descriptor. Only a model with a long-range module, which + the collected metadata names by its topologies, pays for rebuilding the + model. + + Parameters + ---------- + data + Serialized model dictionary, as passed to the export entry point. + metadata + Archive metadata collected from that same model. + + Returns + ------- + Any or None + The evaluated long-range module, or ``None`` when the model folds no + long-range state. + """ + if not metadata.get("topologies"): + return None + + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + model = BaseModel.deserialize(data["model"]) + parts = model.atomic_model.fused_decomposition() + long_range = None if parts is None else parts[0].long_range + if long_range is None or not long_range.charge_state_constants(): + return None + model.to("cpu") + model.eval() + return long_range + + +class _ChargeStateRebuild(torch.nn.Module): + """Rebuild the charge state of a descriptor and of a long-range module. + + The outputs are those of the descriptor fold followed by the folded state + of the long-range module. + + Parameters + ---------- + descriptor + Compressed charge-conditioned descriptor, as returned by + :func:`_charge_state_descriptor`. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`. + """ + + def __init__(self, descriptor: Any, long_range: Any) -> None: + super().__init__() + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + ChargeStateFold, + ) + + self.descriptor = ChargeStateFold(descriptor) + self.long_range = long_range + + def forward(self, charge_spin: torch.Tensor) -> tuple[torch.Tensor, ...]: + return ( + *self.descriptor(charge_spin), + *self.long_range.rebuild_charge_state(charge_spin), + ) + + +def _match_charge_state_constants( + descriptor: Any, exported: Any, long_range: Any | None = None +) -> list[str]: """Name, per fold output, the constant of ``exported`` it replaces. A compressed charge-conditioned descriptor carries its frame condition in @@ -2612,6 +2908,10 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: :func:`_charge_state_descriptor`. exported Exported lower whose constants are to be named. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`, whose state is matched after the + artifacts of the descriptor; ``None`` when there is none. Returns ------- @@ -2632,9 +2932,18 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: for name, tensor in exported.state_dict.items() if isinstance(tensor, torch.Tensor) } + artifacts = [ + (name, getattr(descriptor, f"compress_{name}")) + for name in CHARGE_STATE_ARTIFACTS + ] + if long_range is not None: + artifacts += [ + ("long-range charge state", state) + for state in long_range.charge_state_constants() + ] constants: list[str] = [] - for name in CHARGE_STATE_ARTIFACTS: - artifact = getattr(descriptor, f"compress_{name}").detach().cpu() + for name, tensor in artifacts: + artifact = tensor.detach().cpu() if artifact.numel() == 0: constants.append("") continue @@ -2659,6 +2968,7 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: def _compile_charge_state_fold( descriptor: Any, aoti_configs: dict, + long_range: Any | None = None, ) -> bytes: """Compile the rebuild of the charge-state artifacts as its own archive. @@ -2674,6 +2984,10 @@ def _compile_charge_state_fold( :func:`_charge_state_descriptor`. aoti_configs Inductor options the inference lower was compiled with. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`, whose state the fold rebuilds after + the artifacts of the descriptor; ``None`` when there is none. Returns ------- @@ -2700,7 +3014,12 @@ def _compile_charge_state_fold( dtype=torch.float32, device="cpu", ) - fold = torch.export.export(ChargeStateFold(descriptor), (sample,)) + module = ( + ChargeStateFold(descriptor) + if long_range is None + else _ChargeStateRebuild(descriptor, long_range) + ) + fold = torch.export.export(module, (sample,)) if _env.DEVICE.type != "cpu": from torch.export.passes import ( move_to_device_pass, @@ -2852,11 +3171,12 @@ def _deserialize_to_file_pt2( # state becomes known. charge_state_bytes: bytes | None = None if charge_state_descriptor is not None: + long_range = _charge_state_long_range(data, metadata) metadata["charge_state_constants"] = _match_charge_state_constants( - charge_state_descriptor, exported + charge_state_descriptor, exported, long_range ) charge_state_bytes = _compile_charge_state_fold( - charge_state_descriptor, aoti_configs + charge_state_descriptor, aoti_configs, long_range ) # Second artifact: with-comm. Only for descriptors whose message diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 068bef876f..e88c69fcf1 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -670,6 +670,11 @@ def __init__( type_map: list[str] | None = None, **kwargs: Any, ) -> None: + if kwargs.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) super().__init__( descriptor=descriptor, fitting=fitting_net, type_map=type_map, **kwargs ) diff --git a/deepmd/tf2/atomic_model/dp_atomic_model.py b/deepmd/tf2/atomic_model/dp_atomic_model.py index 2d68ba81ef..966c005b10 100644 --- a/deepmd/tf2/atomic_model/dp_atomic_model.py +++ b/deepmd/tf2/atomic_model/dp_atomic_model.py @@ -41,6 +41,8 @@ class tf2_atomic_model(dpmodel_atomic_model): """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls = None + """No long-range module is implemented in this backend.""" @tf.autograph.experimental.do_not_convert def make_atom_mask( diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 751b6c69b6..86b38996ac 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3549,6 +3549,310 @@ def fitting_variant_type_args() -> Variant: ) +# --- Long-range module configurations: --- # +long_range_args_plugin = ArgsPlugin() +interaction_kernel_args_plugin = ArgsPlugin() + + +@interaction_kernel_args_plugin.register( + "les", + doc="The latent Ewald summation kernel erf(alpha r) / r, the interaction of " + "two Gaussian charge clouds: it reaches 1/r beyond a few 1/alpha and stays " + "finite at contact.", +) +def interaction_kernel_les() -> list[Argument]: + doc_alpha = "The inverse width alpha of the kernel, in Å^-1. Must be positive." + doc_trainable = "Whether alpha is trained." + return [ + Argument("alpha", float, optional=True, default=1.0, doc=doc_alpha), + Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), + ] + + +@interaction_kernel_args_plugin.register( + "sog", + doc="The sum-of-Gaussians kernel sum_k a_k exp(-r^2 / (2 sigma_k^2)).", +) +def interaction_kernel_sog() -> list[Argument]: + doc_amplitude = ( + "The amplitudes a_k of the Gaussians, in Å^-1, of either sign: one value " + "shared by every Gaussian or one value per Gaussian. Required with " + "`width`; with the default widths it defaults to " + "sqrt(2 / pi) ln(`width_ratio`) / sigma_k, the weights of a discretized " + "1 / r, so that the default kernel follows 1 / r, less a nearly " + "constant offset, from a few `base_width` to a fraction of the widest " + "Gaussian. The offset exerts no force between fixed charges; through " + "charges that respond to the positions it does." + ) + doc_width = ( + "The widths sigma_k of the Gaussians, in Å, each positive. Defaults to " + "the geometric sequence `base_width` * `width_ratio`^k for " + "k = 0, ..., `n_gaussian` - 1." + ) + doc_n_gaussian = "The number of Gaussians of the default widths." + doc_base_width = "The width of the first default Gaussian, in Å." + doc_width_ratio = "The ratio of consecutive default widths. Must exceed one." + doc_trainable = "Whether the amplitudes and widths are trained." + return [ + Argument( + "amplitude", + [float, list[float], type(None)], + optional=True, + default=None, + doc=doc_amplitude, + ), + Argument( + "width", + [list[float], type(None)], + optional=True, + default=None, + doc=doc_width, + ), + Argument("n_gaussian", int, optional=True, default=12, doc=doc_n_gaussian), + Argument( + "base_width", + float, + optional=True, + default=2.180230445405648, + doc=doc_base_width, + ), + Argument( + "width_ratio", + float, + optional=True, + default=1.6297670882677647, + doc=doc_width_ratio, + ), + Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), + ] + + +def latent_charge_net_args() -> list[Argument]: + doc_neuron = "The number of neurons in each hidden layer of the charge network." + doc_activation_function = f"The activation function in the charge network. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Defaults to the activation function of the energy fitting." + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection. Defaults to the setting of the energy fitting.' + doc_precision = f"The precision of the charge network and kernel parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Defaults to the precision of the energy fitting." + doc_trainable = "Whether the parameters of the charge network are trained. A list sets every hidden layer and the output layer separately." + doc_seed = "Random seed for parameter initialization of the charge network. Defaults to a seed derived from the seed of the energy fitting." + return [ + Argument("neuron", list[int], optional=True, default=[64, 64], doc=doc_neuron), + Argument( + "activation_function", str, optional=True, doc=doc_activation_function + ), + Argument("resnet_dt", bool, optional=True, doc=doc_resnet_dt), + Argument("precision", str, optional=True, doc=doc_precision), + Argument( + "trainable", + [list[bool], bool], + optional=True, + default=True, + doc=doc_trainable, + ), + Argument("seed", [int, None], optional=True, doc=doc_seed), + ] + + +def charge_constraint_args() -> list[Argument]: + doc_source = ( + "The source of the total charge of a frame: `neutral` for zero, " + "`charge_spin` for the charge of the frame charge/spin condition, which " + "requires a descriptor with `add_chg_spin_ebd`, or `fparam` for the frame " + "parameter at `fparam_index`." + ) + doc_fparam_index = "The column of the frame parameters holding the total charge, used when `source` is `fparam`." + doc_channel = "The latent charge channel that is constrained to the total charge; the other channels stay unconstrained." + return [ + Argument( + "source", + str, + optional=True, + default="neutral", + doc=doc_source, + extra_check=lambda x: x in ("neutral", "charge_spin", "fparam"), + extra_check_errmsg="must be one of 'neutral', 'charge_spin' or 'fparam'", + ), + Argument("fparam_index", int, optional=True, default=0, doc=doc_fparam_index), + Argument("channel", int, optional=True, default=0, doc=doc_channel), + ] + + +@long_range_args_plugin.register( + "latent_charge", + doc="Interaction of latent charges predicted from the descriptor. " + "For a non-periodic system the kernel is summed directly over every pair " + "of atoms; for a periodic system (``periodic: true``) it is summed over " + "the images of the cell in reciprocal space. Every channel interacts " + "through the same kernel; every atom carries half of the interaction of " + "its charges with all others.", +) +def long_range_latent_charge() -> list[Argument]: + doc_channels = "The number of latent charge channels predicted per atom." + doc_charge_net = ( + "The network predicting the latent charges, in units of the elementary " + "charge, from the descriptor; it takes the frame and atomic parameters " + "of the energy fitting." + ) + doc_kernel = "The pair interaction kernel of the latent charges, in Å^-1." + doc_charge_constraint = ( + "Shift one latent charge channel uniformly over the atoms of every frame " + "so that it sums to the total charge of the frame. Unconstrained if not " + "set." + ) + doc_periodic = ( + "Evaluate the module on periodic systems as well: the kernel is summed " + "over the images of the cell in reciprocal space, with the " + "zero-wave-number mode dropped (the tin-foil convention, which places " + "a uniform neutralizing background under a channel with a net " + "charge), less the interaction of every charge with itself. " + "Non-periodic systems keep the direct sum." + ) + doc_kspace_tolerance = ( + "The relative weight of the slowest-decaying kernel mode at the " + "wave-number cutoff k_max of a periodic system." + ) + doc_k_grid_mode = ( + "How training evaluates the reciprocal sum of a periodic batch: " + "'pme', the default of a periodic module, by the smooth particle-mesh " + "Ewald method on a mesh sized by the cells of the batch at " + "`pme_spacing`, whose cost grows nearly linearly with the system; " + "'ragged' over the wave numbers " + "of every frame within k_max, enumerated per batch, so that the k axis " + "is a dynamic dimension; or 'fixed' over one integer grid shared by " + "the frames, sized by the largest cell of the data statistics or by " + "`k_grid_bounds`, with the points a frame does not reach masked, so " + "that the k axis has a static length (a frame whose cell needs points " + "beyond the grid is truncated to it, with a warning). A frozen model " + "evaluates every periodic frame on the mesh of its cell, whatever the " + "mode. 'fixed' and 'pme' require `periodic: true`." + ) + doc_k_grid_bounds = ( + "The per-axis integer bounds [n1, n2, n3] of the fixed k grid, in " + "place of the data statistics; only meaningful with " + "`k_grid_mode: fixed`." + ) + doc_decouple_k_grid = ( + "Decouple the wave-number cutoff of a periodic system from the kernel " + "parameters: it is derived once from the initial parameters and kept " + "for the life of the model, so that the k grid and the mesh stay put " + "while the kernel trains (the kernel values still follow the live " + "parameters). The frozen cutoff is serialized with the model. " + "Requires `periodic: true`." + ) + doc_pme_spacing = ( + "The largest spacing of the mesh of a periodic frame, in Angstrom; " + "the default is three quarters of the Nyquist spacing of the " + "wave-number cutoff, 0.75 pi / k_max. Requires `periodic: true`." + ) + doc_pme_order = ( + "The order of the B-splines that spread the charges to the mesh, 8 " + "by default. Requires `periodic: true`." + ) + return [ + Argument("channels", int, optional=True, default=1, doc=doc_channels), + Argument( + "charge_net", + dict, + latent_charge_net_args(), + optional=True, + doc=doc_charge_net, + ), + Argument( + "kernel", + dict, + [], + [ + Variant( + "type", + interaction_kernel_args_plugin.get_all_argument(), + optional=True, + default_tag="les", + doc="The type of the kernel.", + ) + ], + optional=True, + default={"type": "les"}, + doc=doc_kernel, + ), + Argument( + "charge_constraint", + dict, + charge_constraint_args(), + optional=True, + default=None, + doc=doc_charge_constraint, + ), + Argument("periodic", bool, optional=True, default=False, doc=doc_periodic), + Argument( + "kspace_tolerance", + float, + optional=True, + default=1e-8, + doc=doc_kspace_tolerance, + ), + Argument( + "k_grid_mode", + [str, None], + optional=True, + default=None, + doc=doc_k_grid_mode, + extra_check=lambda x: x in (None, "ragged", "fixed", "pme"), + extra_check_errmsg="must be 'ragged', 'fixed' or 'pme'", + ), + Argument( + "k_grid_bounds", + list[int], + optional=True, + default=None, + doc=doc_k_grid_bounds, + ), + Argument( + "decouple_k_grid", + bool, + optional=True, + default=False, + doc=doc_decouple_k_grid, + ), + Argument( + "pme_spacing", + [float, None], + optional=True, + default=None, + doc=doc_pme_spacing, + ), + Argument( + "pme_order", + [int, None], + optional=True, + default=None, + doc=doc_pme_order, + ), + ] + + +def long_range_args() -> Argument: + doc_long_range = ( + "A long-range module added to the atomic energy of the energy fitting. " + "It reads the same descriptor and is evaluated on complete systems: " + "non-periodic ones directly, periodic ones in reciprocal space when " + "the module serves them. Currently validated with the DPA4C " + "descriptor." + ) + return Argument( + "long_range", + dict, + [], + [ + Variant( + "type", + long_range_args_plugin.get_all_argument(), + doc="The type of the long-range module.", + ) + ], + optional=True, + doc=supported_backends("pt_expt") + doc_long_range, + ) + + # --- Modifier configurations: --- # modifier_args_plugin = ArgsPlugin() @@ -3848,6 +4152,7 @@ def standard_model_args() -> Argument: default={}, doc=supported_backends("pt", "jax", "pd", "pt_expt", "tf2") + doc_info, ), + long_range_args(), *_bridging_method_args(), ], doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index b5bdb2ea94..2fa225cc64 100644 --- a/deepmd/utils/bridging.py +++ b/deepmd/utils/bridging.py @@ -142,6 +142,9 @@ def _is_dpa4_family(sub: dict) -> bool: _LEARNED_CHILD_KEYS = ( "descriptor", "fitting_net", + # A long-range module belongs to the learned energy; the composition then + # rejects it beside the bridging term. + "long_range", "model_branch_alias", "info", "use_compile", diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py index 945805c592..2cf7d82fd4 100644 --- a/deepmd/utils/finetune.py +++ b/deepmd/utils/finetune.py @@ -589,6 +589,18 @@ def _apply_pretrained_model_params( single_config["descriptor"] = deepcopy(pretrained_config["descriptor"]) if not random_fitting: single_config["fitting_net"] = deepcopy(pretrained_config["fitting_net"]) + # A long-range module reads the descriptor beside the fitting and + # is carried over with it; the trainability set in the input + # applies to its parts as to the descriptor and the fitting. + if "long_range" in pretrained_config: + long_range = deepcopy(pretrained_config["long_range"]) + for part in ("charge_net", "kernel"): + long_range.setdefault(part, {})["trainable"] = ( + single_config.get("long_range", {}) + .get(part, {}) + .get("trainable", True) + ) + single_config["long_range"] = long_range log.info( "Change the '%s' model configurations according to pretrained branch '%s'.", model_branch, diff --git a/doc/model/index.rst b/doc/model/index.rst index c6db98d169..b3069c6e7b 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -13,6 +13,7 @@ Model dpa3 dpa4 dpa4c + long-range unimol train-hybrid sel diff --git a/doc/model/long-range.md b/doc/model/long-range.md new file mode 100644 index 0000000000..4c22bc5f96 --- /dev/null +++ b/doc/model/long-range.md @@ -0,0 +1,247 @@ +# Long-range latent charges {{ pytorch_icon }} + +> [!NOTE] +> **Supported backends**: PyTorch Exportable {{ pytorch_icon }} (`dp --pt-expt`) + +A short-range model sees each atom only through the neighbors within its +cutoff. Interactions that decay slowly with distance, such as those between +charged groups of a molecule, therefore escape it. The long-range module adds +such an interaction to the energy of a standard model. A second network reads +the same descriptor as the energy fitting and predicts per-atom _latent +charges_, and every pair of atoms of a system interacts through those charges: + +```math +E = \sum_i E_i^{\text{SR}} + k_e \sum_{i [!WARNING] +> The sum covers complete systems. Without `periodic: true`, a periodic cell +> is rejected. The cost of the non-periodic sum grows with the square of the +> number of atoms of a frame. + +## Configuration + +The module is a `long_range` section of a standard energy model, beside the +descriptor and the fitting. It is currently available with the +[DPA4C](./dpa4c.md) descriptor: + +```json +"model": { + "type_map": ["C", "H", "O"], + "descriptor": { + "type": "dpa4c", + "rcut": 6.0 + }, + "fitting_net": { + "neuron": [128, 128, 128] + }, + "long_range": { + "type": "latent_charge", + "channels": 1, + "kernel": { + "type": "les", + "alpha": 1.0 + }, + "charge_constraint": { + "source": "neutral" + } + } +} +``` + +The charge network, +{ref}`charge_net `, +takes the frame and atomic parameters of the energy fitting and follows its +activation function, skip connection and precision unless they are set; its +seed defaults to one derived from the seed of the fitting. +Its width is set by +{ref}`neuron `. + +### Kernels + +{ref}`kernel ` selects the +interaction of the charges, `les` by default. Both kernels are trained with +the model unless their `trainable` option is false, and their widths stay +positive throughout training. + +- `les`: $K(r) = \mathrm{erf}(\alpha r)/r$, the interaction of two Gaussian + charge clouds. It approaches $1/r$ beyond a few $1/\alpha$ and stays finite, + $2\alpha/\sqrt{\pi}$, at contact. The inverse width $\alpha$ is set by + {ref}`alpha `. +- `sog`: $K(r) = \sum_k a_k \exp(-r^2/2\sigma_k^2)$, a sum of Gaussians whose + amplitudes $a_k$ may take either sign. The widths are given by + {ref}`width ` + or default to a geometric sequence, and the amplitudes by + {ref}`amplitude `, + which explicit widths require. The default amplitudes, + $a_k = \sqrt{2/\pi}\,\ln b/\sigma_k$ for the ratio $b$ of consecutive + widths, discretize $1/r$: with the default widths the kernel stays finite, + $0.46$ Å⁻¹, at contact, and its radial force between fixed charges stays + within 0.4% of that of $1/r$ from 6.5 Å to 150 Å. It differs from $1/r$ + there by a nearly constant offset $C \approx 1.3 \times 10^{-3}$ Å⁻¹, which + exerts no force between fixed charges. The charges respond to the positions, + however, and through them the offset adds + $-k_e C \sum_c (Q_c^2 - \sum_i (q_i^c)^2)/2$ to the energy of a frame, with + $Q_c$ the total charge of channel $c$, which the response of the charges + turns into a force. Below about 5 Å the kernel falls short of $1/r$ by more + than $C$, by 0.12 Å⁻¹ at 2 Å. + +Several charge channels, set by +{ref}`channels `, let +the model represent interactions that one charge per atom cannot. + +### Periodic systems + +With +{ref}`periodic ` enabled, +a frame given a cell is evaluated over the images of the cell. The sum +converges in reciprocal space, where the Fourier transform of the kernel, + +- `les`: $\hat K(\mathbf{k}) = (4\pi/k^2)\, e^{-k^2/4\alpha^2}$, or +- `sog`: $\hat K(\mathbf{k}) = \sum_k a_k (2\pi)^{3/2} \sigma_k^3\, e^{-\sigma_k^2 k^2/2}$, + which with the default amplitudes approaches the Coulomb transform + $4\pi/k^2$, + +decays like a Gaussian. Over the wave numbers $\mathbf{k} = 2\pi B^{-1} +\mathbf{n}$, $\mathbf{n} \in \mathbb{Z}^3 \setminus \{\mathbf{0}\}$, of a cell +$B$ whose rows are the lattice vectors, + +```math +E = k_e \left[\frac{1}{2V}\sum_{\mathbf{k}\neq\mathbf{0}} \sum_{c=1}^{C} +\hat K(\mathbf{k})\, |\hat\rho_c(\mathbf{k})|^2 +- \frac{1}{2} K(0) \sum_{i,c} (q_i^c)^2 \right], +\qquad \hat\rho_c(\mathbf{k}) = \sum_i q_i^c e^{i\mathbf{k}\cdot\mathbf{r}_i}, +``` + +with the cell volume $V$ and the kernel value at contact $K(0)$, whose term +removes the interaction of every charge with itself. The +$\mathbf{k}=\mathbf{0}$ mode is dropped, the tin-foil convention, which places +a uniform neutralizing background under a channel with a net charge. The +wave-number cutoff $k_{max}$ is where the weight of the slowest-decaying mode +falls below +{ref}`kspace_tolerance `, +so the truncation converges exponentially; it follows the kernel as it trains, +unless +{ref}`decouple_k_grid ` +keeps the cutoff of the initial kernel. + +{ref}`k_grid_mode ` +selects how training evaluates the sum: + +- `pme`, the default: by the smooth particle-mesh Ewald method. The charges + are spread with B-splines of order + {ref}`pme_order ` to a + mesh that resolves every cell of the batch at + {ref}`pme_spacing `, + whose transform gives the sum; the cost grows nearly linearly with the + number of atoms. With the default order 8 and spacing $0.75\pi/k_{max}$, the + energy and the forces stay within a few parts in a million of the sum over + the wave numbers. +- `ragged`: over the wave numbers of every frame within $k_{max}$, listed for + every batch; their number grows with the volume of the cell. +- `fixed`: over one integer grid shared by the frames of a batch, sized by the + largest cell of the data statistics or by + {ref}`k_grid_bounds `; + a frame whose cell needs a wave number beyond the grid is truncated to it, + with a warning. + +A frozen model evaluates every periodic frame on the mesh of its cell, +whatever its training mode, so a model trained in another mode can differ +from its frozen form by the accuracy of the mesh. Forces and virials hold the +complete periodic term, and a model with `periodic: true` still evaluates +non-periodic frames by the direct sum. In a `linear_ener` composition, a +sub-model whose module serves periodic systems takes no other sub-model with +a long-range module beside it. + +### Total charge + +{ref}`charge_constraint ` +shifts one channel uniformly over the atoms of every frame so that it sums to +the total charge of the frame. The total charge is zero (`neutral`), the +charge of the frame charge/spin condition of the descriptor (`charge_spin`), +or a column of the frame parameters (`fparam`); a frame without the input +takes its default. The constrained channel then carries the total charge, +while the other channels remain free. Without the section, no channel is +constrained. + +### Exclusions + +Pairs of the types in `pair_exclude_types` see each other neither in the +descriptor nor in the long-range interaction, while the `exclude_types` of +the descriptor act on the descriptor alone. In a `linear_ener` composition, +`pair_exclude_types` belongs to the composition, whose sub-models share its +neighbor graph and pair topology. The reciprocal-space sum of a periodic +system is global and excludes no pair, so `pair_exclude_types` is not +supported with a module that serves periodic systems, nor +`atom_exclude_types` with any long-range module. + +## Training and inference + +Training uses the usual energy loss; see +[training energy models](train-energy.md). With `enable_compile`, the +long-range interaction is part of the compiled model: the pairs, the wave +numbers or the mesh of each batch are built outside the compiled region and +enter it as inputs of dynamic size. + +`dp --pt-expt freeze` exports the model, long-range module included, to +`.pt2` or `.pte`, and `dp --pt-expt compress` compresses its descriptor into a +`.pt2` archive while the networks of the long-range module stay as they are. +A trained checkpoint (`.pt`) and the frozen model are evaluated by `dp test` +and the Python inference interface; the C++ interface and LAMMPS, with +`pair_style deepmd` and `pair_style deepmd/kk`, read the `.pt2` archive. +Since the sum covers the complete system, LAMMPS runs a model whose module +serves no periodic system with non-periodic boundaries (`f`, `s` or `m`) in +every dimension, and a model with `periodic: true` with boundaries periodic in +every dimension or in none. Several models given to `pair_style deepmd` for +the model deviation evaluate a non-periodic system together only when every +one of them carries a long-range module. A model that the dedicated operator described +below evaluates, a compressed DPA4C model frozen for a CUDA device, also runs +on several MPI ranks: every rank gathers the positions and charges of the +atoms of all ranks over MPI and evaluates its own atoms against all of them, +so that the ranks share the cost of the direct sum, while each evaluates the +mesh of a periodic system whole. This requires DeePMD-kit built with MPI. Any +other model, like an uncompressed one, runs on a single rank; +`pair_style deepmd/kk` evaluates an uncompressed model on non-periodic systems +only. + +Compressing a descriptor conditioned on `charge_spin` specializes the model +to one charge state at a time, and a constraint that reads its total charge +from `charge_spin` follows the same state. The compressed archive serves any +charge state: the Python interface applies the state of each call, and LAMMPS +the one given by the `charge_spin` keyword of `pair_style deepmd`. + +A compressed DPA4C model frozen for a CUDA device evaluates the long-range +term with a dedicated operator, which reads the positions of the atoms and +the cell. It enumerates the pairs of a non-periodic system itself, so that its +memory grows linearly with the number of atoms and its cost with their +square, and evaluates a periodic system on the mesh of its cell. The operator +serves the `les` kernel and the `sog` kernel of up to 32 Gaussians, with up to +8 charge channels, B-splines of order up to 12, a charge network without +frame or atomic parameters, and no `pair_exclude_types`. Any other model, +like an uncompressed one, receives the pairs or the mesh of each frame from +the interface, so the memory of a non-periodic evaluation grows with the +square of the number of atoms. + +The following are not supported for a model with a long-range module: +multi-rank (domain-decomposed) inference of a model that the dedicated +operator does not evaluate, such as an uncompressed one, spin models, the +tabulated ZBL potential, and short-range bridging. diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index cf40990302..2bfe93694c 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -457,13 +457,20 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, * @param[in] world Pointer to the MPI communicator of the caller (an * ``MPI_Comm``), as the ``world`` argument of DP_NewNlist_comm, when the * system spans several ranks, or NULL on a single rank. A model with a - * long-range module evaluates a complete non-periodic system: on a single - * rank its nodes are the whole system (``nall_nodes == nloc``); on several - * ranks it requires the communicator, reads the positions of the nodes each - * rank owns, and is evaluated on every rank at every step, a rank without - * nodes included (``nall_nodes``, ``nloc`` and ``edge_storage`` zero), whose - * buffers sized by the nodes or by the edge storage may then be NULL. Every - * other model ignores it. + * long-range module evaluates a complete system: a single rank holds every + * node of a non-periodic system (``nall_nodes == nloc``), and the halo nodes + * of a periodic one are images of its owned nodes; on several ranks it + * requires the communicator, reads the positions of the nodes each rank owns, + * and is evaluated on every rank at every step, a rank without nodes included + * (``nall_nodes``, ``nloc`` and ``edge_storage`` zero), whose buffers sized by + * the nodes or by the edge storage may then be NULL. Every other model + * ignores it. + * @param[in] box The cell of a periodic system in host memory, nine values in + * the length unit of the model, row-major with the lattice vectors as rows, + * or NULL for a non-periodic system. A model with a long-range module that + * serves periodic systems evaluates the system over the images of the cell, + * one whose module serves none rejects a cell, and every other model ignores + * it. * @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. @@ -473,9 +480,9 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, * without a synchronization of the stream. A graph without halo nodes is * captured into a CUDA graph on the second consecutive call with the same * buffers and extents and replayed while they stay the same, except where a - * model with a long-range module receives a communicator. The host waits for - * the device when a call captures or releases a graph, which happens on that - * second call and on a call whose buffers or extents change. + * model with a long-range module receives a communicator or a cell. The host + * waits for the device when a call captures or releases a graph, which + * happens on that second call and on a call whose buffers or extents change. * DP_DeepPotSetChargeSpin waits for the work in flight on the model device * before it replaces the constants and for the replacement before it returns, * and the caller completes the work on its stream before it deletes the model. @@ -487,8 +494,8 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, * under way invalidates it, and the model then evaluates without replay. * @note API versions 28 to 30 named this function * DP_DeepPotComputeCanonicalGraphGPU, without ``d_virial``, ``d_coord``, - * ``world`` and ``stream``; a binary built against that name fails to - * load rather than call a function of another argument list. + * ``world``, ``box`` and ``stream``; a binary built against that name + * fails to load rather than call a function of another argument list. * @since API version 31 */ extern void DP_DPA4ComputeCanonicalGraphGPU( @@ -506,6 +513,7 @@ extern void DP_DPA4ComputeCanonicalGraphGPU( const int64_t* d_source_order, const double* d_coord, void* world, + const double* box, int64_t nloc, int64_t nall_nodes, int64_t edge_storage, @@ -539,10 +547,11 @@ extern bool DP_DeepPotCanonicalGraphRequiresSourceCSR(DP_DeepPot* dp); /** * @brief Query whether the model carries a long-range module. * - * Such a module sums its interaction over every pair of atoms, so the model - * evaluates only a complete non-periodic system, and the device-resident - * entry DP_DPA4ComputeCanonicalGraphGPU passes it the node positions and, on - * several ranks, the communicator. + * Such a module sums its interaction over every pair of atoms, and over the + * images of the cell of a periodic system when it serves periodic systems, so + * the model evaluates complete systems, and the device-resident entry + * DP_DPA4ComputeCanonicalGraphGPU passes it the node positions, the cell and, + * on several ranks, the communicator. * @since API version 31 */ extern bool DP_DeepPotHasLongRange(DP_DeepPot* dp); diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 90925416b8..7fe1560b05 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1363,9 +1363,10 @@ class DeepPot : public DeepBaseModel { * ``d_virial`` may be null when the virial is not requested. * ``d_atom_virial`` may be null, and then so may the source CSR unless * canonical_graph_requires_source_csr() holds. ``d_coord`` holds the node - * positions and ``world`` the MPI communicator of a system on several - * ranks, which a model with a long-range module reads (has_long_range()); - * both may be null for any other model. A non-null ``stream`` is the + * positions, ``world`` the MPI communicator of a system on several ranks + * and ``box`` the cell of a periodic system, which a model with a + * long-range module reads (has_long_range()); all three may be null for + * any other model. A non-null ``stream`` is the * caller's accelerator stream on the model device. The arguments have the * semantics of DP_DPA4ComputeCanonicalGraphGPU. */ @@ -1382,6 +1383,7 @@ class DeepPot : public DeepBaseModel { const int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const int64_t nloc, const int64_t nall_nodes, const int64_t edge_storage, @@ -1389,7 +1391,8 @@ class DeepPot : public DeepBaseModel { DP_DPA4ComputeCanonicalGraphGPU( dp, d_atom_energy, d_force, d_virial, d_atom_virial, d_atype, d_source, index_width, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, - d_source_order, d_coord, world, nloc, nall_nodes, edge_storage, stream); + d_source_order, d_coord, world, box, nloc, nall_nodes, edge_storage, + stream); DP_CHECK_OK(DP_DeepPotCheckOK, dp); } @@ -1405,8 +1408,7 @@ class DeepPot : public DeepBaseModel { /** * @brief Query whether the model carries a long-range module, which - * evaluates a complete non-periodic system only; see - * DP_DeepPotHasLongRange. + * evaluates complete systems; see DP_DeepPotHasLongRange. */ bool has_long_range() const { const bool result = DP_DeepPotHasLongRange(dp); diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index 01b5a25b9e..eeb196fead 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -1831,16 +1831,17 @@ void DP_DPA4ComputeCanonicalGraphGPU(DP_DeepPot* dp, const int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const int64_t nloc, const int64_t nall_nodes, const int64_t edge_storage, void* stream) { - DP_REQUIRES_OK(dp, - dp->dp.compute_canonical_graph_gpu( - d_atom_energy, d_force, d_virial, d_atom_virial, d_atype, - d_source, index_width, d_edge_vec, d_destination_row_ptr, - d_source_row_ptr, d_source_order, d_coord, world, nloc, - nall_nodes, edge_storage, stream)); + DP_REQUIRES_OK( + dp, dp->dp.compute_canonical_graph_gpu( + d_atom_energy, d_force, d_virial, d_atom_virial, d_atype, + d_source, index_width, d_edge_vec, d_destination_row_ptr, + d_source_row_ptr, d_source_order, d_coord, world, box, nloc, + nall_nodes, edge_storage, stream)); } bool DP_DeepPotSupportsDeviceEdgeInference(DP_DeepPot* dp) { diff --git a/source/api_c/tests/test_deeppot_long_range_hpp.cc b/source/api_c/tests/test_deeppot_long_range_hpp.cc new file mode 100644 index 0000000000..9e6804e056 --- /dev/null +++ b/source/api_c/tests/test_deeppot_long_range_hpp.cc @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test the C API with a pt_expt (.pt2) DPA4C model that carries a +// latent-charge long-range module (source/tests/infer/gen_dpa4c_long_range.py). +#include + +#include +#include +#include + +#include "../../api_cc/tests/expected_ref.h" +#include "DeepPotPTExpt.h" +#include "deepmd.hpp" +#include "test_utils.h" + +namespace { +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4c_long_range.pt2"; +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4c_long_range.expected"; +constexpr const char* kShortRangeModelPath = + "../../tests/infer/deeppot_sea.pth"; +} // namespace + +template +class TestInferDeepPotLongRangeHPP : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector expected_f; + double expected_tot_e = 0.; + deepmd::hpp::DeepPot dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + dp.init(kModelPath); + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_f = ref.get("default", "expected_f"); + for (const VALUETYPE value : ref.get("default", "expected_e")) { + expected_tot_e += value; + } + } +}; + +TYPED_TEST_SUITE(TestInferDeepPotLongRangeHPP, ValueTypes); + +TYPED_TEST(TestInferDeepPotLongRangeHPP, has_long_range) { + EXPECT_TRUE(this->dp.has_long_range()); +} + +TEST(TestInferDeepPotModelDeviLongRangeHPP, has_long_range) { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + // A long-range model anywhere in the ensemble confines all of them. + deepmd::hpp::DeepPotModelDevi mixed; + mixed.init(std::vector{kShortRangeModelPath, kModelPath}); + EXPECT_TRUE(mixed.has_long_range()); + deepmd::hpp::DeepPotModelDevi short_range; + short_range.init( + std::vector{kShortRangeModelPath, kShortRangeModelPath}); + EXPECT_FALSE(short_range.has_long_range()); +} + +TYPED_TEST(TestInferDeepPotLongRangeHPP, nopbc) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector()); + ASSERT_EQ(force.size(), this->expected_f.size()); + EXPECT_LT(fabs(ener - this->expected_tot_e), epsilon); + for (std::size_t ii = 0; ii < force.size(); ++ii) { + EXPECT_LT(fabs(force[ii] - this->expected_f[ii]), epsilon); + } +} diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index f47af3bea7..a7fcdc6479 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -436,6 +436,7 @@ class DeepPotBackend : public DeepBaseModelBackend { const std::int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const std::int64_t nloc, const std::int64_t nall_nodes, const std::int64_t edge_storage, @@ -915,9 +916,10 @@ class DeepPot : public DeepBaseModel { * @brief Whether the model carries a long-range module. * * Such a module sums its interaction over every pair of atoms of the - * system, so the model evaluates only a complete non-periodic system, and - * compute_canonical_graph_gpu() passes it the positions of the nodes and, - * on several ranks, the communicator. + * system, and over the images of the cell of a periodic system when it + * serves periodic systems, so the model evaluates complete systems, and + * compute_canonical_graph_gpu() passes it the positions of the nodes, the + * cell and, on several ranks, the communicator. */ bool has_long_range() const; @@ -948,13 +950,20 @@ class DeepPot : public DeepBaseModel { * @param[in] world Pointer to the MPI communicator of the caller (an * ``MPI_Comm``), as InputNlist::world, when the system spans several * ranks, or null on a single rank. A model with a long-range module - * evaluates a complete non-periodic system: on a single rank its nodes are - * the whole system (``nall_nodes == nloc``); on several ranks it requires + * evaluates a complete system: a single rank holds every node of a + * non-periodic system (``nall_nodes == nloc``), and the halo nodes of a + * periodic one are images of its owned nodes; on several ranks it requires * the communicator, reads the positions of the nodes each rank owns, and * is evaluated on every rank at every step, a rank without nodes included * (``nall_nodes``, ``nloc`` and ``edge_storage`` zero), whose buffers * sized by the nodes or by the edge storage may then be null. Every other * model ignores it. + * @param[in] box The cell of a periodic system in host memory, nine values + * in the length unit of the model, row-major with the lattice vectors as + * rows, or null for a non-periodic system. A model with a long-range + * module that serves periodic systems evaluates the system over the images + * of the cell, one whose module serves none rejects a cell, and every + * other model ignores it. * @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. @@ -964,11 +973,11 @@ class DeepPot : public DeepBaseModel { * enqueued on it without a synchronization of the stream. A graph without * halo nodes is captured into a CUDA graph on the second consecutive call * with the same buffers and extents and replayed while they stay the same, - * except where a model with a long-range module receives a communicator. - * The host waits for the device when a call captures or releases a graph, - * which happens on that second call and on a call whose buffers or extents - * change. set_charge_spin waits for the work in flight on the model device - * before it replaces the constants and for the replacement before it + * except where a model with a long-range module receives a communicator or + * a cell. The host waits for the device when a call captures or releases a + * graph, which happens on that second call and on a call whose buffers or + * extents change. set_charge_spin waits for the work in flight on the model + * device before it replaces the constants and for the replacement before it * returns, and the caller completes the work on its stream before it * destroys the model. A null stream runs on the current stream of the * calling thread for the model device and waits for that stream before @@ -991,6 +1000,7 @@ class DeepPot : public DeepBaseModel { const std::int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const std::int64_t nloc, const std::int64_t nall_nodes, const std::int64_t edge_storage, @@ -1233,5 +1243,14 @@ class DeepPotModelDevi : public DeepBaseModelDevi { protected: std::vector> dps; + + private: + /** + * @brief Check that every model accepts the cell handed to the ensemble. + * Without one, the neighbor-list entries serve a model with a long-range + * module beside models that take the cell of the region. + **/ + template + void check_ensemble_box(const std::vector& dbox) const; }; } // namespace deepmd diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 9cd0569673..71e66a2109 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -389,6 +389,7 @@ class DeepPotPTExpt : public DeepPotBackend { const std::int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const std::int64_t nloc, const std::int64_t nall_nodes, const std::int64_t edge_storage, @@ -422,6 +423,7 @@ class DeepPotPTExpt : public DeepPotBackend { bool uses_fp32_edge_vectors() const override; bool uses_canonical_graph_inference() const override; bool canonical_graph_requires_source_csr() const override; + bool has_long_range() const override; private: template @@ -452,6 +454,7 @@ class DeepPotPTExpt : public DeepPotBackend { const std::int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const std::int64_t nloc, const std::int64_t nall_nodes, const std::int64_t edge_storage, @@ -494,6 +497,15 @@ class DeepPotPTExpt : public DeepPotBackend { at::Tensor source_row_placeholder_; at::Tensor source_order_placeholder_; bool graph_edge_fp32_ = false; + // The model carries a long-range module, whose lower reads the topologies + // of the frame after its other inputs: the pair lists, and the mesh of the + // cell of a module that serves periodic systems, for the graph lower; the + // node positions, the communicator of the ranks and the cell for the compact + // canonical lower. + bool long_range_ = false; + // The module serves periodic systems, on a mesh of at most this spacing. + bool long_range_periodic_ = false; + double mesh_spacing_ = 0.0; NeighborListData nlist_data; at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) std::vector mapping_; // cached mapping vector (LAMMPS path) @@ -645,6 +657,11 @@ class DeepPotPTExpt : public DeepPotBackend { * @param[in] destination_row_ptr Destination CSR offsets ``(N + 1,)``. * @param[in] source_order Source-grouped edge permutation ``(E,)``. * @param[in] source_row_ptr Source CSR offsets ``(N + 1,)``. + * @param[in] coord Positions of the nodes of the single frame, shape + * ``(N, 3)`` float64, from which the topologies of a model with a + * long-range module are built; read only for such a model. + * @param[in] box The row-major cell of a periodic frame, nine values, or + * empty for a non-periodic one; read only for such a model. */ std::vector run_model_graph( const torch::Tensor& atype, @@ -659,8 +676,22 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& source_row_ptr, const torch::Tensor& fparam, const torch::Tensor& aparam, - const torch::Tensor& charge_spin); + const torch::Tensor& charge_spin, + const torch::Tensor& coord, + const std::vector& box); + /** + * @brief Run a compact canonical ``.pt2`` artifact. + * + * @param[in] coord Node positions, shape ``(N, 3)`` float64, appended for a + * model with a long-range module and ignored otherwise. + * @param[in] world The MPI communicator (a pointer to an ``MPI_Comm``) of the + * ranks a decomposed system spans, or NULL for a system that one rank holds + * whole; a model with a long-range module receives it after the positions. + * @param[in] box The row-major cell of a periodic frame, nine values, or + * empty for a non-periodic one; a model with a long-range module receives + * it after the communicator. + */ std::vector run_model_canonical_graph( const torch::Tensor& atype, const torch::Tensor& n_node, @@ -669,7 +700,30 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& edge_vec, const torch::Tensor& destination_row_ptr, const torch::Tensor& source_row_ptr, - const torch::Tensor& source_order); + const torch::Tensor& source_order, + const torch::Tensor& coord, + const void* world, + const std::vector& box); + + /** + * @brief Evaluate the long-range module on a rank that holds no node. + * + * The compact lower runs on one halo node without edges, which carries no + * charge and contributes nothing, so that the rank joins the exchanges of + * the module with the other ranks of ``world``, in the cell ``box`` of the + * system, empty for a non-periodic one. + */ + void join_long_range_exchange(const void* world, + const std::vector& box); + + /** + * @brief Check the cell given to a model with a long-range module. + * + * @param[in] box The row-major cell of a periodic frame, nine values, or + * empty for a non-periodic one, which every long-range module serves; a + * cell requires a module that serves periodic systems. + */ + void check_long_range_box(const std::vector& box) const; /** * @brief Run the with-comm .pt2 artifact with comm tensors appended. diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 3f0e889ee2..68de144850 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -670,6 +670,7 @@ void DeepPotBackend::compute_canonical_graph_gpu( const std::int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const std::int64_t nloc, const std::int64_t nall_nodes, const std::int64_t edge_storage, @@ -686,6 +687,7 @@ void DeepPotBackend::compute_canonical_graph_gpu( (void)d_source_order; (void)d_coord; (void)world; + (void)box; (void)nloc; (void)nall_nodes; (void)edge_storage; @@ -774,6 +776,7 @@ void DeepPot::compute_canonical_graph_gpu( const std::int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const std::int64_t nloc, const std::int64_t nall_nodes, const std::int64_t edge_storage, @@ -781,7 +784,8 @@ void DeepPot::compute_canonical_graph_gpu( dp->compute_canonical_graph_gpu( d_atom_energy, d_force, d_virial, d_atom_virial, d_atype, d_source, index_width, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, - d_source_order, d_coord, world, nloc, nall_nodes, edge_storage, stream); + d_source_order, d_coord, world, box, nloc, nall_nodes, edge_storage, + stream); } bool DeepPot::uses_fp32_edge_vectors() const { @@ -947,6 +951,25 @@ template void DeepPotModelDevi::compute( const std::vector& aparam, const std::vector& charge_spin); +// LAMMPS hands a model with a long-range module no cell for a non-periodic +// system, whereas the neighbor-list entries of other models, such as those of +// the TensorFlow backend, read the cell of the region. +template +void DeepPotModelDevi::check_ensemble_box( + const std::vector& dbox) const { + if (!dbox.empty() || !has_long_range()) { + return; + } + for (const auto& model : dps) { + if (!model->has_long_range()) { + throw deepmd::deepmd_exception( + "an ensemble that holds a model with a long-range module evaluates " + "a non-periodic system only when every model carries one; use " + "models with long-range modules throughout, or a periodic system."); + } + } +} + template void DeepPotModelDevi::compute(std::vector& all_energy, std::vector>& all_force, @@ -963,6 +986,7 @@ void DeepPotModelDevi::compute(std::vector& all_energy, if (numb_models == 0) { return; } + check_ensemble_box(dbox); all_energy.resize(numb_models); all_force.resize(numb_models); all_virial.resize(numb_models); @@ -1025,6 +1049,7 @@ void DeepPotModelDevi::compute( all_virial.resize(numb_models); all_atom_energy.resize(numb_models); all_atom_virial.resize(numb_models); + check_ensemble_box(dbox); for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_virial[ii], all_atom_energy[ii], all_atom_virial[ii], dcoord_, datype_, diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 37a991f8ad..e6ec9d9dd5 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -22,6 +23,7 @@ #include "commonPTExpt.h" #include "device.h" #include "errors.h" +#include "longRangePTExpt.h" #include "neighbor_list.h" using deepmd::ptexpt::check_call_charge_spin; @@ -119,6 +121,28 @@ FrameParameterLayout resolve_frame_parameter_layout(const char* name, " for frame-specific values."); } +// The communicator input of the compact lower of a long-range model: an int64 +// CPU tensor holding the address of the ``MPI_Comm`` that ``world`` points to, +// or zero for a system that one rank holds whole. The handle itself is read +// only by the operator, with the width of ``MPI_Comm`` in its MPI library. +at::Tensor communicator_tensor(const void* world) { + return torch::full( + {1}, static_cast(reinterpret_cast(world)), + torch::TensorOptions().dtype(torch::kInt64)); +} + +// The cell input of the compact lower of a long-range model: a float64 CPU +// tensor of shape (1, 9), the rows of the row-major cell its lattice vectors, +// zero for a non-periodic frame. +at::Tensor cell_tensor(const std::vector& box) { + auto cell = + torch::zeros({1, 9}, torch::TensorOptions().dtype(torch::kFloat64)); + if (!box.empty()) { + std::copy(box.begin(), box.begin() + 9, cell.data_ptr()); + } + return cell; +} + } // namespace void DeepPotPTExpt::translate_error(std::function f) { @@ -267,6 +291,33 @@ void DeepPotPTExpt::init(const std::string& model, lower_input_is_canonical_ = false; canonical_requires_source_csr_ = false; } + // A model with a long-range module reads the all-pairs topology of the + // frame after its other inputs, and, when the module serves periodic + // systems, the mesh of the cell, sized by the spacing of the metadata. + long_range_ = metadata.obj_val.count("topologies") > 0; + long_range_periodic_ = false; + if (long_range_) { + for (const auto& name : metadata["topologies"].as_array()) { + if (name.as_string() == "pme_mesh") { + long_range_periodic_ = true; + } else if (name.as_string() != "all_pairs") { + throw deepmd::deepmd_exception("unknown topology '" + name.as_string() + + "' in the .pt2 metadata; expected " + "'all_pairs' or 'pme_mesh'."); + } + } + if (long_range_periodic_) { + mesh_spacing_ = + metadata["topology_params"]["pme_mesh"]["spacing"].as_double(); + } + if (!lower_input_is_graph_ && !lower_input_is_canonical_) { + throw deepmd::deepmd_exception( + "a model with a long-range module requires the graph or the " + "compact canonical lower."); + } + // Its compact lower assembles the force through the source CSR. + canonical_requires_source_csr_ = lower_input_is_canonical_; + } if (lower_input_is_edge_) { std::cerr << "WARNING: This .pt2 uses the deprecated edge_vec lower " "schema (pt-backend SeZM/DPA4 freeze). Support will be " @@ -550,7 +601,9 @@ std::vector DeepPotPTExpt::run_model_graph( const torch::Tensor& source_row_ptr, const torch::Tensor& fparam, const torch::Tensor& aparam, - const torch::Tensor& charge_spin) { + const torch::Tensor& charge_spin, + const torch::Tensor& coord, + const std::vector& box) { // Graph-input ABI: original edge payload plus destination/source CSR views. check_graph_aparam_flat(aparam, daparam, "run_model_graph"); std::vector inputs = { @@ -566,6 +619,21 @@ std::vector DeepPotPTExpt::run_model_graph( if (dchgspin > 0) { inputs.push_back(charge_spin); } + if (long_range_) { + // A periodic frame holds no pair; the mesh describes its interactions. + const auto vec_dtype = graph_edge_fp32_ ? torch::kFloat32 : torch::kFloat64; + const auto pairs = + box.empty() + ? deepmd::ptexpt::buildAllPairsTopology( + coord, atype, pair_exclude_table_, ntypes, vec_dtype) + : deepmd::ptexpt::buildGuardPairTopology(coord.device(), vec_dtype); + inputs.insert(inputs.end(), pairs.begin(), pairs.end()); + if (long_range_periodic_) { + const auto mesh = deepmd::ptexpt::buildMeshTopology( + coord, box, mesh_spacing_, vec_dtype); + inputs.insert(inputs.end(), mesh.begin(), mesh.end()); + } + } const std::lock_guard evaluation(evaluation_mutex_); return loader->run(inputs); } @@ -578,11 +646,71 @@ std::vector DeepPotPTExpt::run_model_canonical_graph( const torch::Tensor& edge_vec, const torch::Tensor& destination_row_ptr, const torch::Tensor& source_row_ptr, - const torch::Tensor& source_order) { + const torch::Tensor& source_order, + const torch::Tensor& coord, + const void* world, + const std::vector& box) { + std::vector inputs = { + atype, n_node, + n_local, canonicalIndexBytes(source), + edge_vec, destination_row_ptr, + source_row_ptr, source_order}; + if (long_range_) { + inputs.push_back(coord); + inputs.push_back(communicator_tensor(world)); + inputs.push_back(cell_tensor(box)); + } const std::lock_guard evaluation(evaluation_mutex_); - return loader->run({atype, n_node, n_local, canonicalIndexBytes(source), - edge_vec, destination_row_ptr, source_row_ptr, - source_order}); + return loader->run(inputs); +} + +void DeepPotPTExpt::check_long_range_box(const std::vector& box) const { + if (!box.empty() && box.size() != 9) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates one frame with a cell of " + "nine values, or without one, but got " + + std::to_string(box.size()) + " values."); + } + if (!box.empty() && !long_range_periodic_) { + throw deepmd::deepmd_exception( + "the long-range module of this model serves complete non-periodic " + "systems only; pass no cell (in LAMMPS, a non-periodic boundary f, s " + "or m in every dimension), or train the module with periodic: true."); + } + if (!box.empty()) { + // A cell whose volume is a vanishing fraction of the product of the + // lengths of its lattice vectors, or not finite, spans no lattice. + const double volume = + std::abs(box[0] * (box[4] * box[8] - box[5] * box[7]) - + box[1] * (box[3] * box[8] - box[5] * box[6]) + + box[2] * (box[3] * box[7] - box[4] * box[6])); + double lengths = 1.0; + for (int d = 0; d < 3; ++d) { + lengths *= + std::sqrt(box[3 * d] * box[3 * d] + box[3 * d + 1] * box[3 * d + 1] + + box[3 * d + 2] * box[3 * d + 2]); + } + if (!(volume > 1e-6 * lengths)) { + throw deepmd::deepmd_exception( + "the cell of a periodic system needs three independent lattice " + "vectors, but its volume is " + + std::to_string(volume) + "; pass no cell for a non-periodic system."); + } + } +} + +void DeepPotPTExpt::join_long_range_exchange(const void* world, + const std::vector& box) { + const torch::Device device(torch::kCUDA, gpu_id); + const auto i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); + // The compact lower stores at least two edge slots; without edges both are + // padding outside every CSR row. + run_model_canonical_graph( + torch::zeros({1}, i64), torch::ones({1}, i64), torch::zeros({1}, i64), + torch::zeros({2}, i64.dtype(torch::kInt32)), + torch::zeros({2, 3}, i64.dtype(torch::kFloat32)), torch::zeros({2}, i64), + torch::zeros({2}, i64), torch::arange(2, i64), + torch::zeros({1, 3}, i64.dtype(torch::kFloat64)), world, box); } std::vector DeepPotPTExpt::run_model_with_comm( @@ -839,6 +967,40 @@ 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; + // A long-range module sums over every pair of atoms of a non-periodic + // system, which the caller marks with an empty box, and over the images of + // the cell of a periodic one. A single rank holds every atom, so a ghost + // atom there is a periodic image. Across ranks, the compact lower gathers + // the owned atoms of every rank through the communicator; the pairs that + // the graph lower lists stay within one. + const std::vector long_range_box(box.begin(), box.end()); + if (long_range_) { + check_long_range_box(long_range_box); + } + if (long_range_ && long_range_box.empty() && !multi_rank && nghost_real > 0) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a complete non-periodic " + "system without a box, but this single-rank call carries " + + std::to_string(nghost_real) + + " ghost atoms, the periodic images of a periodic cell; pass the cell " + "of a periodic system to a model whose module serves one."); + } + if (long_range_ && multi_rank && !lower_input_is_canonical_) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a system decomposed over " + + std::to_string(lmp_list.nprocs) + + " MPI ranks only through the compact canonical lower, which a " + "compressed DPA4C model frozen for a CUDA device uses when the " + "long-range operator serves its long-range module; this archive uses " + "another lower, so run it on a single rank."); + } + if (long_range_ && multi_rank && lmp_list.world == nullptr) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a system decomposed over " + "several MPI ranks through their communicator, but the neighbor list " + "carries none (InputNlist::world)."); + } + const void* const world = multi_rank ? lmp_list.world : nullptr; // Whether the edge topology folds ghost neighbours onto their local owners, // which reads an owner for every extended atom out of ``mapping_``. The // edge lower folds unless the with-comm artifact carries ghost features @@ -1267,7 +1429,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // ``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.) + // ghosts within rcut -- so it still runs the model normally.) A + // long-range module across ranks still joins the exchanges of the + // others. + if (long_range_ && multi_rank) { + join_long_range_exchange(world, long_range_box); + } ener.assign(nframes, static_cast(0)); force.assign(static_cast(nframes) * fwd_map.size() * 3, static_cast(0)); @@ -1308,19 +1475,22 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, graph_pack.atype = node_atype; graph_pack.n_node = n_node_tensor; graph_pack.n_local = n_local_tensor; + const at::Tensor node_coord = + coord_Tensor.slice(1, 0, n_node_count).reshape({n_node_count, 3}); 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); + compact.source_row_ptr, compact.source_order, node_coord, world, + long_range_box); } else { flat_outputs = run_model_graph( node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask, graph_pack.destination_order, graph_pack.destination_row_ptr, graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, - graph_aparam, charge_spin_tensor); + graph_aparam, charge_spin_tensor, node_coord, long_range_box); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -1513,6 +1683,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, check_call_charge_spin(charge_spin, /*nframes=*/1, settable_chgspin, dchgspin > 0, default_chg_spin_, chg_spin_table_ranges_); + // A long-range module sums over the atoms of a non-periodic system, which + // the empty box marks, and over the images of the cell of a periodic one. + const std::vector long_range_box(box.begin(), box.end()); + if (long_range_) { + check_long_range_box(long_range_box); + } // The .pt2 model only contains forward_common_lower, which requires // nlist as input. We must build the nlist in C++ and fold back the // extended-region outputs to local atoms. @@ -1688,12 +1864,14 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, if (graph_edge_fp32_) { graph_tensors.edge_vec = graph_tensors.edge_vec.to(torch::kFloat32); } + const at::Tensor node_coord = + coord_Tensor.slice(1, 0, nloc).reshape({nloc, 3}); 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); + compact.source_order, node_coord, /*world=*/nullptr, long_range_box); } else { flat_outputs = run_model_graph( graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, @@ -1705,7 +1883,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // normalizes the rectangular runtime aparam to the flat // (N, daparam) node axis of the graph ABI. extend_graph_aparam(aparam_tensor, natoms, natoms, daparam), - charge_spin_tensor); + charge_spin_tensor, node_coord, long_range_box); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -2261,6 +2439,15 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, 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 (long_range_ && (extended || comm_nlist != nullptr)) { + throw deepmd::deepmd_exception( + "compute_edges_gpu: a model with a long-range module evaluates a " + "complete non-periodic system on a single rank (nall_nodes == nloc and " + "no comm_nlist); periodic systems and several ranks require the " + "compact canonical lower, which a compressed DPA4C model frozen for a " + "CUDA device uses when the long-range operator serves its module, or " + "the evaluation from a neighbor list (pair_style deepmd)."); + } if (lower_input_is_graph_ && extended && has_message_passing_) { throw deepmd::deepmd_exception( "compute_edges_gpu: message-passing graph-input inference requires an " @@ -2488,7 +2675,8 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask, graph_pack.destination_order, graph_pack.destination_row_ptr, graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, - graph_aparam, charge_spin_tensor)); + graph_aparam, charge_spin_tensor, coord_t.reshape({nloc, 3}), + /*box=*/{})); 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(); @@ -2586,6 +2774,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( const std::int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const std::int64_t nloc, const std::int64_t nall_nodes, const std::int64_t edge_storage, @@ -2598,6 +2787,39 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( throw deepmd::deepmd_exception( "compute_canonical_graph_gpu requires a CUDA device."); } + // A long-range module reads the positions of a complete system and the + // cell of a periodic one: the node set of a single rank, whose halo nodes + // are images of a periodic cell, or the owned nodes of every rank of + // ``world``, which a rank without nodes still joins. + const std::vector long_range_box = + box != nullptr ? std::vector(box, box + 9) + : std::vector(); + if (long_range_) { + check_long_range_box(long_range_box); + } + if (long_range_ && d_coord == nullptr && nall_nodes > 0) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu: a model with a long-range module reads " + "the node positions (d_coord)."); + } + if (long_range_ && world == nullptr && box == nullptr && nall_nodes != nloc) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu: a model with a long-range module on a " + "single rank evaluates a complete non-periodic system " + "(nall_nodes == nloc), or a periodic one with its cell; a system " + "decomposed over several ranks takes their communicator."); + } + if (long_range_ && world != nullptr && nall_nodes == 0) { + translate_error([&] { + const c10::DeviceGuard device_guard(torch::Device(torch::kCUDA, gpu_id)); + const ptexpt::CallerStreamScope stream_scope(stream, gpu_id); + join_long_range_exchange(world, long_range_box); + if (stream == nullptr) { + synchronize_current_accelerator_stream(gpu_id); + } + }); + return; + } if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2) { throw deepmd::deepmd_exception( "invalid compact canonical graph dimensions."); @@ -2681,11 +2903,15 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( // captured graph owns them. auto n_node = torch::full({1}, nall_nodes, opt_i64); auto n_local = torch::full({1}, nloc, opt_i64); + const auto coord = d_coord != nullptr + ? torch::from_blob(const_cast(d_coord), + {nall_nodes, 3}, opt_f64) + : torch::Tensor(); 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)); + extract_outputs(output, run_model_canonical_graph( + atype, n_node, n_local, source, edge_vec, + destination_row_ptr, source_row_ptr, + source_order, coord, world, long_range_box)); auto atom_energy = output["atom_energy"] .reshape({nall_nodes}) .slice(0, 0, nloc) @@ -2718,9 +2944,11 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( // neighbor rebuild, so a graph with halo nodes is evaluated eagerly. A // model with a long-range module reads the node positions, whose address // is then part of the key, and is evaluated eagerly over several ranks, - // where it may communicate between its kernels. + // where it may communicate between its kernels, and in a cell, which the + // evaluation reads on the host. const bool long_range = has_long_range(); - if (nloc != nall_nodes || (long_range && world != nullptr)) { + if (nloc != nall_nodes || + (long_range && (world != nullptr || box != nullptr))) { evaluate(); } else { if (!canonical_replay_) { @@ -2809,6 +3037,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( const std::int64_t* d_source_order, const double* d_coord, void* world, + const double* box, const std::int64_t nloc, const std::int64_t nall_nodes, const std::int64_t edge_storage, @@ -2816,7 +3045,8 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( compute_canonical_graph_gpu_impl( d_atom_energy, d_force, d_virial, d_atom_virial, d_atype, d_source, index_width, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, - d_source_order, d_coord, world, nloc, nall_nodes, edge_storage, stream); + d_source_order, d_coord, world, box, nloc, nall_nodes, edge_storage, + stream); } bool DeepPotPTExpt::uses_fp32_edge_vectors() const { @@ -2837,4 +3067,6 @@ bool DeepPotPTExpt::canonical_graph_requires_source_csr() const { return canonical_requires_source_csr_; } +bool DeepPotPTExpt::has_long_range() const { return long_range_; } + #endif diff --git a/source/api_cc/src/longRangePTExpt.h b/source/api_cc/src/longRangePTExpt.h new file mode 100644 index 0000000000..352f85a643 --- /dev/null +++ b/source/api_cc/src/longRangePTExpt.h @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include + +#include +#include +#include +#include + +#include "commonPT.h" + +namespace deepmd { +namespace ptexpt { + +/** + * @brief Build the all-pairs topology a long-range module of the graph lower + * reads. + * + * Inference-path twin of Python ``build_pair_graph`` in + * ``deepmd/dpmodel/utils/neighbor_graph/pair_graph.py`` for a single frame: + * every unordered pair of nodes once, ``[src, dst]`` with ``src < dst``, its + * vector ``r_src - r_dst``, and two masked guard pairs that keep the pair axis + * at the lower bound of the exported lower. The model-level pair exclusion is + * applied to the mask, as it is to the neighbor graph. + * + * @param coord Node positions, shape ``(N, 3)``, float64, on the model device. + * @param atype Node types, shape ``(N,)``, int64. + * @param type_mask_table Device keep table of the pair exclusion, or an + * undefined tensor when the model excludes no pair. + * @param ntypes Number of atom types. + * @param vec_dtype Precision of the pair vectors, that of the edge vectors of + * the artifact. + * @return ``{pair_index (2, P + 2) int64, pair_vec (P + 2, 3), pair_mask + * (P + 2,) bool}``. + */ +inline std::vector buildAllPairsTopology( + const torch::Tensor& coord, + const torch::Tensor& atype, + const torch::Tensor& type_mask_table, + const int ntypes, + const torch::ScalarType vec_dtype) { + const std::int64_t node_count = coord.size(0); + const auto index_options = + torch::TensorOptions().dtype(torch::kInt64).device(coord.device()); + const auto mask_options = + torch::TensorOptions().dtype(torch::kBool).device(coord.device()); + const torch::Tensor index = + torch::triu_indices(node_count, node_count, 1, index_options); + const torch::Tensor vec = (coord.index_select(0, index.select(0, 0)) - + coord.index_select(0, index.select(0, 1))) + .to(vec_dtype); + const torch::Tensor mask = + applyPairExclusion(index, torch::ones({index.size(1)}, mask_options), + atype, type_mask_table, ntypes); + return {torch::cat({index, torch::zeros({2, 2}, index_options)}, 1), + torch::cat({vec, torch::zeros({2, 3}, vec.options())}, 0), + torch::cat({mask, torch::zeros({2}, mask_options)})}; +} + +/** + * @brief Build the all-pairs topology of a periodic frame, which holds no + * pair: the reciprocal-space topology describes the interactions. + * + * Counterpart of Python ``empty_pair_graph`` for the exported lower: two + * masked guard pairs, the lower bound of its pair axis. + */ +inline std::vector buildGuardPairTopology( + const torch::Device& device, const torch::ScalarType vec_dtype) { + const auto options = torch::TensorOptions().device(device); + return {torch::zeros({2, 2}, options.dtype(torch::kInt64)), + torch::zeros({2, 3}, options.dtype(vec_dtype)), + torch::zeros({2}, options.dtype(torch::kBool))}; +} + +/** + * @brief The smallest even integer of at least ``n`` without a prime factor + * above 7. + */ +inline std::int64_t evenFftSize(std::int64_t n) { + std::int64_t m = std::max(2, n + n % 2); + while (true) { + std::int64_t rest = m; + for (const std::int64_t p : {2, 3, 5, 7}) { + while (rest % p == 0) { + rest /= p; + } + } + if (rest == 1) { + return m; + } + m += 2; + } +} + +/** + * @brief Build the mesh topology a periodic long-range module of the graph + * lower reads. + * + * Twin of Python ``build_mesh_grid`` in + * ``deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py`` for a single + * frame: along every lattice vector ``a_d`` the mesh holds the smallest even + * number of points without a prime factor above 7 of at least + * ``ceil(|a_d| / spacing)``, with the length formed term by term as Python + * forms it, so that both size the mesh identically. Without a cell, the frame + * is non-periodic: the cell is the identity, the mesh minimal and the frame + * masked. + * + * @param coord Node positions, shape ``(N, 3)``, on the model device. + * @param box The row-major cell, the rows its lattice vectors, nine values in + * the length unit of the model, or empty for a non-periodic frame. + * @param spacing The largest mesh spacing, in the length unit of the model. + * @param dtype Precision of the geometry, that of the edge vectors. + * @return ``{box (1, 3, 3), volume (1,), coord (N, 3), strain (1, 3, 3), + * mask (1,) bool, mode_1 (n1,), mode_2 (n2,), mode_3 (n3,)}``. + */ +inline std::vector buildMeshTopology( + const torch::Tensor& coord, + const std::vector& box, + const double spacing, + const torch::ScalarType dtype) { + const auto options = + torch::TensorOptions().dtype(dtype).device(coord.device()); + const bool periodic = !box.empty(); + std::array cell = {1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0}; + std::array dims = {2, 2, 2}; + if (periodic) { + for (int k = 0; k < 9; ++k) { + cell[k] = box[k]; + } + for (int d = 0; d < 3; ++d) { + const double* v = cell.data() + 3 * d; + const double length = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + dims[d] = + evenFftSize(static_cast(std::ceil(length / spacing))); + } + } + const double volume = + std::abs(cell[0] * (cell[4] * cell[8] - cell[5] * cell[7]) - + cell[1] * (cell[3] * cell[8] - cell[5] * cell[6]) + + cell[2] * (cell[3] * cell[7] - cell[4] * cell[6])); + // The copy detaches the tensor from the local array, which a conversion + // to float64 on the host would alias. + const auto cell_tensor = + torch::from_blob(cell.data(), {1, 3, 3}, torch::kFloat64) + .to(options, /*non_blocking=*/false, /*copy=*/true); + std::vector fields = { + cell_tensor, torch::full({1}, volume, options), coord.to(dtype), + torch::zeros({1, 3, 3}, options), + torch::full({1}, periodic, options.dtype(torch::kBool))}; + for (const std::int64_t n : dims) { + // The integer frequencies of an axis of n points, in FFT order. + const auto m = torch::arange(n, options); + fields.push_back(torch::where(m > n / 2, m - n, m)); + } + return fields; +} + +} // namespace ptexpt +} // namespace deepmd diff --git a/source/api_cc/tests/dpa4c_canonical_test_common.h b/source/api_cc/tests/dpa4c_canonical_test_common.h index 882cb4e7ad..4d6fa22d81 100644 --- a/source/api_cc/tests/dpa4c_canonical_test_common.h +++ b/source/api_cc/tests/dpa4c_canonical_test_common.h @@ -537,8 +537,8 @@ class TestInferDpa4cCanonicalPtExpt : public ::testing::Test { atom_energy, force, total_virial, atom_virial, atype, source, deepmd::canonical_index_width(nodes), edge_vec, destination_row_ptr, source_row_ptr, source_order, - /*d_coord=*/nullptr, /*world=*/nullptr, natoms, nodes, edge_storage, - stream); + /*d_coord=*/nullptr, /*world=*/nullptr, /*box=*/nullptr, natoms, + nodes, edge_storage, stream); } } diff --git a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc new file mode 100644 index 0000000000..b45ea775b3 --- /dev/null +++ b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc @@ -0,0 +1,603 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test C++ inference for pt_expt (.pt2) DPA4C models with a latent-charge +// long-range module, which sums over every pair of atoms of a complete +// non-periodic system, and over the images of the cell of a periodic one when +// the module serves periodic systems. Models generated by +// source/tests/infer/gen_dpa4c_long_range.py: uncompressed archives on the +// graph lower, which receives the pair list or the mesh of the cell, and +// compressed archives on the compact canonical lower, which reads the +// positions of the atoms and the cell. +#include + +#include +#include +#include + +#include "DeepPot.h" +#include "DeepPotPTExpt.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT +#include + +#include +#include +#include + +#include "commonPT.h" +#endif + +namespace { +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4c_long_range.pt2"; +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4c_long_range.expected"; +constexpr const char* kCompressedModelPath = + "../../tests/infer/deeppot_dpa4c_long_range_compressed.pt2"; +constexpr const char* kCompressedRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_compressed.expected"; +constexpr const char* kPeriodicModelPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic.pt2"; +constexpr const char* kPeriodicRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic.expected"; +constexpr const char* kPeriodicCompressedModelPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic_compressed.pt2"; +constexpr const char* kPeriodicCompressedRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic_compressed.expected"; + +// The full neighbor list of a non-periodic system within ``rc``: every atom +// is local, so the list holds no ghost. +template +std::vector> nopbc_nlist(const std::vector& coord, + const double rc) { + const int natoms = static_cast(coord.size() / 3); + std::vector> nlist(natoms); + for (int ii = 0; ii < natoms; ++ii) { + for (int jj = 0; jj < natoms; ++jj) { + double r2 = 0.; + for (int dd = 0; dd < 3; ++dd) { + const double diff = coord[ii * 3 + dd] - coord[jj * 3 + dd]; + r2 += diff * diff; + } + if (jj != ii && r2 < rc * rc) { + nlist[ii].push_back(jj); + } + } + } + return nlist; +} +} // namespace + +// One reference case of an archive: per-atom energies, forces and virials, +// with the totals the interface returns. +template +struct LongRangeRef { + std::vector expected_f; + double expected_tot_e = 0.; + std::vector expected_tot_v = std::vector(9, 0.); + + void load(const deepmd_test::ExpectedRef& ref, const char* section) { + const std::vector atom_e = + ref.template get(section, "expected_e"); + const std::vector atom_v = + ref.template get(section, "expected_v"); + expected_f = ref.template get(section, "expected_f"); + for (std::size_t ii = 0; ii < atom_e.size(); ++ii) { + expected_tot_e += atom_e[ii]; + for (int dd = 0; dd < 9; ++dd) { + expected_tot_v[dd] += atom_v[ii * 9 + dd]; + } + } + } + + void check(double ener, + const std::vector& force, + const std::vector& virial, + double epsilon) const { + ASSERT_EQ(force.size(), expected_f.size()); + ASSERT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), epsilon); + for (std::size_t ii = 0; ii < force.size(); ++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); + } + } +}; + +template +class TestInferDeepPotDpa4cLongRangePtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector charge_spin_explicit = {1.0, 2.0}; + LongRangeRef ref_default; + LongRangeRef ref_explicit; + + static deepmd::DeepPot dp; + + static void SetUpTestSuite() { +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT + dp.init(kModelPath); +#endif + } + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + ref_default.load(ref, "default"); + ref_explicit.load(ref, "explicit"); + } + + static void TearDownTestSuite() { dp = deepmd::DeepPot(); } + + // Computes through the LAMMPS-style entry point, with the full neighbor + // list of the system and no ghost atom. + void compute_lmp(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& charge_spin) { + std::vector> nlist_data = nopbc_nlist(coord, dp.cutoff()); + const int nloc = static_cast(atype.size()); + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, virial, coord, atype, std::vector(), 0, + inlist, 0, {}, {}, charge_spin); + } +}; + +template +deepmd::DeepPot TestInferDeepPotDpa4cLongRangePtExpt::dp; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangePtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, has_long_range) { + EXPECT_TRUE(this->dp.has_long_range()); + EXPECT_EQ(this->dp.dim_chg_spin(), 2); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, standalone_default) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector()); + this->ref_default.check(ener, force, virial, epsilon); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, standalone_explicit) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector(), {}, {}, + this->charge_spin_explicit); + this->ref_explicit.check(ener, force, virial, epsilon); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, lmp_nlist_explicit) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->compute_lmp(ener, force, virial, this->charge_spin_explicit); + this->ref_explicit.check(ener, force, virial, epsilon); +} + +// A module that serves no periodic system rejects a cell. +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, periodic_box_is_rejected) { + using VALUETYPE = TypeParam; + const std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + double ener; + std::vector force, virial; + EXPECT_THROW( + this->dp.compute(ener, force, virial, this->coord, this->atype, box), + deepmd::deepmd_exception); +} + +// The compressed archive folds the charge state into constants, which +// set_charge_spin rewrites: the tables of the descriptor and the total charge +// of the constraint together. +template +class TestInferDeepPotDpa4cLongRangeCompressedPtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + LongRangeRef ref_default; + LongRangeRef ref_explicit; + deepmd::DeepPot dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(kCompressedModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kCompressedModelPath + << " was not generated (gen_dpa4c_long_range.py writes it " + "on a CUDA device only)."; + } + dp.init(kCompressedModelPath); + deepmd_test::ExpectedRef ref; + ref.load(kCompressedRefPath); + ref_default.load(ref, "default"); + ref_explicit.load(ref, "explicit"); + } + + void compute_lmp(double& ener, + std::vector& force, + std::vector& virial) { + std::vector> nlist_data = nopbc_nlist(coord, dp.cutoff()); + const int nloc = static_cast(atype.size()); + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, virial, coord, atype, std::vector(), 0, + inlist, 0); + } +}; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangeCompressedPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4cLongRangeCompressedPtExpt, + set_charge_spin_follows_the_state) { + using VALUETYPE = TypeParam; + constexpr double epsilon = 1e-4; + EXPECT_TRUE(this->dp.has_long_range()); + double ener; + std::vector force, virial; + this->compute_lmp(ener, force, virial); + this->ref_default.check(ener, force, virial, epsilon); + this->dp.set_charge_spin({1.0, 2.0}); + this->compute_lmp(ener, force, virial); + this->ref_explicit.check(ener, force, virial, epsilon); + this->dp.set_charge_spin({0.0, 1.0}); + this->compute_lmp(ener, force, virial); + this->ref_default.check(ener, force, virial, epsilon); +} + +// A module that serves periodic systems evaluates a periodic system over the +// images of its cell, through the coordinates and the cell or through the +// neighbor list of a LAMMPS-style call, whose ghost atoms are images, and +// keeps the direct sum of a system without a cell. ``Archive`` selects the +// uncompressed archive on the graph lower or, on a CUDA device, the +// compressed one on the compact canonical lower. +template +class TestInferDeepPotDpa4cLongRangePeriodicPtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13.0, 0.0, 0.0, 0.6, 12.5, + 0.0, -0.4, 0.8, 13.5}; + LongRangeRef ref_periodic; + LongRangeRef ref_cluster; + deepmd::DeepPot dp; + + void load(const char* model_path, const char* ref_path) { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(model_path); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << model_path + << " was not generated (gen_dpa4c_long_range.py writes the " + "compressed archives on a CUDA device only)."; + } + dp.init(model_path); + deepmd_test::ExpectedRef ref; + ref.load(ref_path); + ref_periodic.load(ref, "periodic"); + ref_cluster.load(ref, "cluster"); + } + + // Computes the periodic system through the LAMMPS-style entry point, with + // the ghost atoms of the cell, and folds their forces back onto their + // owners. + void compute_lmp_periodic(double& ener, + std::vector& force, + std::vector& virial) { + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector> nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, static_cast(dp.cutoff())); + const int nloc = static_cast(atype.size()); + const int nall = static_cast(atype_cpy.size()); + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + inlist.mapping = mapping.data(); + std::vector force_all; + dp.compute(ener, force_all, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0); + _fold_back(force, force_all, mapping, nloc, nall, 3); + } + + void check_all(const double epsilon) { + EXPECT_TRUE(dp.has_long_range()); + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box); + ref_periodic.check(ener, force, virial, epsilon); + compute_lmp_periodic(ener, force, virial); + ref_periodic.check(ener, force, virial, epsilon); + dp.compute(ener, force, virial, coord, atype, std::vector()); + ref_cluster.check(ener, force, virial, epsilon); + } +}; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, graph_lower) { + using VALUETYPE = TypeParam; + this->load(kPeriodicModelPath, kPeriodicRefPath); + if (this->IsSkipped()) { + return; + } + this->check_all(std::is_same::value ? 1e-8 : 1e-4); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, compact_lower) { + this->load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath); + if (this->IsSkipped()) { + return; + } + this->check_all(1e-4); +} + +// A cell without volume marks no system: it is rejected rather than taken as +// a non-periodic one. +TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, + degenerate_cell_is_rejected) { + using VALUETYPE = TypeParam; + this->load(kPeriodicModelPath, kPeriodicRefPath); + if (this->IsSkipped()) { + return; + } + double ener; + std::vector force, virial; + EXPECT_THROW(this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector(9, 0.)), + deepmd::deepmd_exception); +} + +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT +// The device-resident entry point of the compressed archives, which reads the +// positions of the nodes and the cell of a periodic system besides the +// compact canonical graph. A non-periodic system that one rank holds whole is +// evaluated eagerly, captured into a CUDA graph and replayed; a periodic one +// is evaluated eagerly, on the local atoms with the ghost atoms folded onto +// their owners, as one rank builds the graph, or with a node for every ghost +// atom, an image of its owner. Every evaluation writes into outputs reset to +// NaN, so an unwritten output fails the comparison. +class TestInferDeepPotDpa4cLongRangeDevicePtExpt : public ::testing::Test { + protected: + // A replay whose capture fails evaluates eagerly instead, with the same + // results; only this warning tells the two apart. + static constexpr const char* kCaptureWarning = + "could not capture the evaluation"; + + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13.0, 0.0, 0.0, 0.6, 12.5, 0.0, -0.4, 0.8, 13.5}; + deepmd::DeepPot dp; + deepmd_test::ExpectedRef ref; + bool capturing_stderr = false; + + // A compact canonical graph on the device with the positions of its nodes, + // the atom that owns every node, and outputs reused by its evaluations, + // whose addresses the replay of repeated evaluations keys on. + struct Graph { + deepmd::CanonicalGraphTensorPack tensors; + torch::Tensor coord; + std::vector owner; + int nodes = 0; + torch::Tensor atom_energy; + torch::Tensor force; + torch::Tensor virial; + torch::Tensor atom_virial; + }; + + // The outputs of an evaluation on the host, the forces folded onto the + // owners of the nodes. + struct Results { + std::vector atom_energy; + std::vector force; + std::vector virial; + }; + + void load(const char* model_path, const char* ref_path) { + if (!torch::cuda::is_available()) { + GTEST_SKIP() << "Skip because the compact canonical path needs a CUDA " + "device."; + } + std::ifstream model_file(model_path); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << model_path + << " was not generated (gen_dpa4c_long_range.py writes the " + "compressed archives on a CUDA device only)."; + } + dp.init(model_path); + ASSERT_TRUE(dp.uses_canonical_graph_inference()); + ASSERT_TRUE(dp.has_long_range()); + ASSERT_TRUE(dp.canonical_graph_requires_source_csr()); + ref.load(ref_path); + testing::internal::CaptureStderr(); + capturing_stderr = true; + } + + void TearDown() override { + if (!capturing_stderr) { + return; + } + const std::string log = testing::internal::GetCapturedStderr(); + std::cerr << log; + EXPECT_EQ(log.find(kCaptureWarning), std::string::npos) + << "an evaluation fell back from the replay to eager evaluation"; + } + + // The graph of the neighbor list ``nlist`` of the extended atoms + // ``coord_ext``, of which the first ``nloc`` are local: folded onto the + // local atoms, or with a node for every extended atom. + Graph build_graph(const std::vector>& nlist, + const std::vector& coord_ext, + const std::vector& atype_ext, + const std::vector& mapping, + const int nloc, + const bool fold) const { + const torch::Device device(torch::kCUDA, 0); + const int nall = static_cast(atype_ext.size()); + const std::vector mapping_64(mapping.begin(), mapping.end()); + deepmd::GraphTensorPack pack = deepmd::buildGraphTensors( + nlist, coord_ext, atype_ext, mapping_64, nloc, nall, dp.cutoff(), + device, /*row_centers=*/nullptr, fold); + deepmd::canonicalizeGraphPayload(pack, pack.atype.size(0)); + pack.edge_vec = pack.edge_vec.to(torch::kFloat32); + Graph graph; + graph.tensors = deepmd::compactCanonicalGraph(pack); + deepmd::CanonicalGraphTensorPack& compact = graph.tensors; + compact.atype = compact.atype.to(torch::kInt64).contiguous(); + compact.source = compact.source.contiguous(); + compact.edge_vec = compact.edge_vec.contiguous(); + compact.destination_row_ptr = compact.destination_row_ptr.contiguous(); + compact.source_row_ptr = compact.source_row_ptr.contiguous(); + compact.source_order = compact.source_order.contiguous(); + graph.nodes = static_cast(compact.atype.size(0)); + // The nodes are the leading extended atoms, in their order. + std::vector positions( + coord_ext.begin(), coord_ext.begin() + 3 * std::size_t(graph.nodes)); + const auto f64 = torch::TensorOptions().dtype(torch::kFloat64); + graph.coord = torch::from_blob(positions.data(), {graph.nodes, 3}, f64) + .to(device) + .contiguous(); + graph.owner = mapping; + graph.owner.resize(graph.nodes); + const auto f64_device = f64.device(device); + graph.atom_energy = torch::empty({nloc}, f64_device); + graph.force = torch::empty({graph.nodes, 3}, f64_device); + graph.virial = torch::empty({9}, f64_device); + graph.atom_virial = torch::empty({graph.nodes, 9}, f64_device); + return graph; + } + + // Evaluates ``graph`` through compute_canonical_graph_gpu on the current + // stream, in the cell ``cell`` (null for a non-periodic system). + Results evaluate(Graph& graph, const int nloc, const double* cell) { + const double nan = std::numeric_limits::quiet_NaN(); + graph.atom_energy.fill_(nan); + graph.force.fill_(nan); + graph.virial.fill_(nan); + graph.atom_virial.fill_(nan); + const deepmd::CanonicalGraphTensorPack& compact = graph.tensors; + dp.compute_canonical_graph_gpu( + graph.atom_energy.data_ptr(), graph.force.data_ptr(), + graph.virial.data_ptr(), graph.atom_virial.data_ptr(), + compact.atype.data_ptr(), compact.source.data_ptr(), + deepmd::canonical_index_width(graph.nodes), + compact.edge_vec.data_ptr(), + compact.destination_row_ptr.data_ptr(), + compact.source_row_ptr.data_ptr(), + compact.source_order.data_ptr(), + graph.coord.data_ptr(), /*world=*/nullptr, cell, nloc, + graph.nodes, compact.source.size(0)); + const auto host_vector = [](const torch::Tensor& tensor) { + const torch::Tensor values = tensor.cpu().contiguous(); + return std::vector(values.data_ptr(), + values.data_ptr() + values.numel()); + }; + Results results; + results.atom_energy = host_vector(graph.atom_energy); + results.virial = host_vector(graph.virial); + const std::vector force = host_vector(graph.force); + results.force.assign(3 * std::size_t(nloc), 0.0); + for (int node = 0; node < graph.nodes; ++node) { + for (int dd = 0; dd < 3; ++dd) { + results.force[3 * graph.owner[node] + dd] += force[3 * node + dd]; + } + } + return results; + } + + // Agreement of an evaluation with the reference of a section. + void expect_reference(const Results& results, const char* section) const { + LongRangeRef expected; + expected.load(ref, section); + double energy = 0.0; + for (const double value : results.atom_energy) { + energy += value; + } + expected.check(energy, results.force, results.virial, 1e-4); + } +}; + +TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, cluster_is_replayed) { + ASSERT_NO_FATAL_FAILURE(load(kCompressedModelPath, kCompressedRefPath)); + if (IsSkipped()) { + return; + } + const int nloc = static_cast(atype.size()); + std::vector mapping(nloc); + for (int ii = 0; ii < nloc; ++ii) { + mapping[ii] = ii; + } + Graph graph = build_graph(nopbc_nlist(coord, dp.cutoff()), coord, atype, + mapping, nloc, /*fold=*/true); + ASSERT_EQ(graph.nodes, nloc); + // Eager, captured and replayed. + const Results eager = evaluate(graph, nloc, nullptr); + expect_reference(eager, "default"); + for (int call = 1; call < 3; ++call) { + const Results replayed = evaluate(graph, nloc, nullptr); + EXPECT_EQ(replayed.atom_energy, eager.atom_energy) << "call " << call; + EXPECT_EQ(replayed.force, eager.force) << "call " << call; + EXPECT_EQ(replayed.virial, eager.virial) << "call " << call; + } + // The positions are part of the key of the replay: the same values at + // another address are evaluated anew. + graph.coord = graph.coord.clone(); + const Results moved = evaluate(graph, nloc, nullptr); + EXPECT_EQ(moved.atom_energy, eager.atom_energy); + EXPECT_EQ(moved.force, eager.force); +} + +TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, periodic_cell) { + ASSERT_NO_FATAL_FAILURE( + load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath)); + if (IsSkipped()) { + return; + } + std::vector> nlist; + std::vector coord_ext; + std::vector atype_ext, mapping; + _build_nlist(nlist, coord_ext, atype_ext, mapping, coord, atype, box, + static_cast(dp.cutoff())); + const int nloc = static_cast(atype.size()); + for (const bool fold : {true, false}) { + Graph graph = build_graph(nlist, coord_ext, atype_ext, mapping, nloc, fold); + ASSERT_EQ(graph.nodes, fold ? nloc : static_cast(atype_ext.size())); + for (int call = 0; call < 2; ++call) { + expect_reference(evaluate(graph, nloc, box.data()), "periodic"); + } + } +} +#endif diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index d875ae618c..58913df34f 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -161,6 +161,12 @@ else: env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_canonical.py & PID16=$! fi + # DPA4C with a latent-charge long-range module: the uncompressed archive + # takes the all-pairs topology of the graph lower; for the CUDA variant the + # script also writes the compressed archive, whose compact canonical lower + # reads the node positions. + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_long_range.py & + PID16=$! wait $PID11 wait $PID12 wait $PID13 @@ -170,6 +176,7 @@ else: wait $PID15 wait $PID16 fi + wait $PID16 fi if [ "${ENABLE_PADDLE:-TRUE}" == "TRUE" ]; then PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 1b559afe50..2ed86f2da2 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -178,6 +178,11 @@ double PairDeepMD::eval_energy_with_fparam( dbox[7] = domain->h[3] / dist_unit_cvt_factor; // zy dbox[6] = domain->h[4] / dist_unit_cvt_factor; // zx dbox[3] = domain->h[5] / dist_unit_cvt_factor; // yx + // A long-range module reads the cell of a periodic system only; the empty + // box marks a non-periodic one. + if (long_range && domain->nonperiodic) { + dbox.clear(); + } // get coord for (int ii = 0; ii < nall; ++ii) { @@ -297,6 +302,11 @@ void PairDeepMD::compute(int eflag, int vflag) { dbox[7] = domain->h[3] / dist_unit_cvt_factor; // zy dbox[6] = domain->h[4] / dist_unit_cvt_factor; // zx dbox[3] = domain->h[5] / dist_unit_cvt_factor; // yx + // A long-range module reads the cell of a periodic system only; the empty + // box marks a non-periodic one. + if (long_range && domain->nonperiodic) { + dbox.clear(); + } // get coord for (int ii = 0; ii < nall; ++ii) { @@ -966,6 +976,24 @@ void PairDeepMD::settings(int narg, char** arg) { all_force.resize(numb_models); } +void PairDeepMD::init_style() { + PairDeepBaseModel::init_style(); + // A long-range module sums over every pair of atoms of a non-periodic + // system, or over the images of the cell of a system periodic in every + // dimension; the atoms of the other ranks reach it through the + // communicator. The model checks which of the two its module serves. + long_range = numb_models > 1 ? deep_pot_model_devi.has_long_range() + : deep_pot.has_long_range(); + const int periodic_dimensions = + domain->xperiodic + domain->yperiodic + domain->zperiodic; + if (long_range && periodic_dimensions != 0 && periodic_dimensions != 3) { + error->all(FLERR, + "pair style deepmd: a model with a long-range module evaluates " + "a system periodic in every dimension or in none; use p in " + "every dimension, or f, s or m in every dimension."); + } +} + /* ---------------------------------------------------------------------- set coeffs for one or more type pairs ------------------------------------------------------------------------- */ @@ -1082,24 +1110,6 @@ void PairDeepMD::coeff(int narg, char** arg) { /* ---------------------------------------------------------------------- */ -void PairDeepMD::init_style() { - PairDeepBaseModel::init_style(); - // A long-range module sums over every pair of atoms of the system, to which - // a periodic box would add the periodic images. - const bool long_range = numb_models > 1 ? deep_pot_model_devi.has_long_range() - : deep_pot.has_long_range(); - if (long_range && - (domain->xperiodic || domain->yperiodic || domain->zperiodic)) { - error->all(FLERR, - "pair style {}: a model with a long-range module evaluates a " - "complete non-periodic system; use a non-periodic boundary (f, " - "s or m) in every dimension.", - force->pair_style); - } -} - -/* ---------------------------------------------------------------------- */ - int PairDeepMD::pack_reverse_comm(int n, int first, double* buf) { int i, m, last; diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h index ac8881b553..69466f740a 100644 --- a/source/lmp/pair_deepmd.h +++ b/source/lmp/pair_deepmd.h @@ -59,6 +59,9 @@ class PairDeepMD : public PairDeepBaseModel { // 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(); + // The model carries a long-range module, which reads the cell of a + // periodic system and no cell for a non-periodic one. + bool long_range = false; private: CommBrickDeepMD* commdata_; diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index ebecf39e6f..2d24f2c4e9 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -2,6 +2,7 @@ #ifdef LMP_KOKKOS #include "pair_deepmd_kokkos.h" +#include #include #include #include @@ -35,7 +36,6 @@ PairDeepMDKokkos::PairDeepMDKokkos(LAMMPS* lmp) edge_vec_fp32(false), canonical_graph(false), canonical_source_csr(false), - long_range(false), device_path_ok(false), reverse_virial(false), reverse_used_host(false) { @@ -223,8 +223,15 @@ void PairDeepMDKokkos::init_style() { canonical_graph = deep_pot.uses_canonical_graph_inference(); canonical_source_csr = deep_pot.canonical_graph_requires_source_csr(); // PairDeepMD::init_style() has confined a model with a long-range module to - // a complete non-periodic system. - long_range = deep_pot.has_long_range(); + // a system periodic in every dimension or in none. The device edge path + // passes such a module no cell, so a periodic system reaches it through the + // compact canonical lower alone. + if (long_range && !canonical_graph && !domain->nonperiodic) { + error->all(FLERR, + "pair style deepmd/kk evaluates a periodic system with a " + "long-range module only through a compressed DPA4C model frozen " + "for a CUDA device; use pair_style deepmd for this model."); + } // 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. @@ -407,10 +414,15 @@ template const double* PairDeepMDKokkos::model_node_positions() { atomKK->sync(execution_space, X_MASK); auto x = atomKK->k_x.template view(); + // Relative to the origin of the box, as the plain pair style passes them. // Without virtual atoms the model nodes are the leading atoms, whose - // coordinates are already in the length unit of the model at a conversion - // factor of one. - if (!compact_graph.has_null_types && dist_unit_cvt_factor == 1.0) { + // coordinates are already those positions in the length unit of the model + // for a box at the origin and a conversion factor of one. + const double x0 = domain->boxlo[0]; + const double y0 = domain->boxlo[1]; + const double z0 = domain->boxlo[2]; + if (!compact_graph.has_null_types && dist_unit_cvt_factor == 1.0 && + x0 == 0.0 && y0 == 0.0 && z0 == 0.0) { return x.data(); } const int nnode_model = compact_graph.nnode_model; @@ -425,9 +437,9 @@ const double* PairDeepMDKokkos::model_node_positions() { "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) * inv_dist; - coord_model(3 * m + 1) = x(i, 1) * inv_dist; - coord_model(3 * m + 2) = x(i, 2) * inv_dist; + coord_model(3 * m + 0) = (x(i, 0) - x0) * inv_dist; + coord_model(3 * m + 1) = (x(i, 1) - y0) * inv_dist; + coord_model(3 * m + 2) = (x(i, 2) - z0) * inv_dist; }); return coord_model.data(); } @@ -574,6 +586,17 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // stream that the model shares. const double* coord_ptr = (!canonical_graph || long_range) ? model_node_positions() : nullptr; + // A long-range module reads the cell of a periodic system, row-major with + // the lattice vectors as rows, in the length unit of the model, and no + // cell for a non-periodic one. + const double* h = domain->h; + std::array cell = {h[0], 0.0, 0.0, h[5], h[1], + 0.0, h[4], h[3], h[2]}; + for (double& entry : cell) { + entry /= dist_unit_cvt_factor; + } + const double* cell_ptr = + long_range && !domain->nonperiodic ? cell.data() : nullptr; void* const model_stream = compact_canonical_stream(typename DeviceType::execution_space()); if (!canonical_graph || model_stream == nullptr) { @@ -589,8 +612,8 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { compact_graph.d_destination_row_ptr.data(), source_csr ? compact_graph.d_source_row_ptr.data() : nullptr, source_csr ? compact_graph.d_source_order.data() : nullptr, - coord_ptr, world_ptr, nloc_m, nnode_m, compact_graph.storage_count, - model_stream); + coord_ptr, world_ptr, cell_ptr, nloc_m, nnode_m, + compact_graph.storage_count, model_stream); } else if (edge_vec_fp32) { deep_pot.compute_edges_gpu(d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), coord_ptr, diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index f65e112109..5858a0895f 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -83,8 +83,9 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // neighbor list, returning the edge count. Public because it launches // extended device lambdas, which CUDA forbids inside non-public members. int build_edges_device(); - // Positions of the model nodes, in model-node order and in the length unit - // of the model; public for the same reason. + // Positions of the model nodes relative to the origin of the box, in + // model-node order and in the length unit of the model; public for the same + // reason. const double* model_node_positions(); protected: @@ -132,7 +133,6 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // The canonical artifact reads the source CSR on every call, not only for // the per-atom virial. bool canonical_source_csr; - bool long_range; // the model reads the positions of its nodes bool device_path_ok; // resolved once in init_style bool reverse_virial; // reverse communication operates on centroid virial bool reverse_used_host; // the host pack/unpack overrides ran diff --git a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py index a574cef821..a73908d6f3 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py +++ b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py @@ -51,6 +51,13 @@ "domain decomposition (nswap>0). Pass '1 1 1' for a single-rank " "reference run on the same archive (single-artifact dispatch).", ) +parser.add_argument( + "--boundary", + type=str, + default="p p p", + help="LAMMPS boundary style. Default 'p p p'; a model with a long-range " + "module evaluates complete non-periodic systems and passes 'f f f'.", +) parser.add_argument( "--pair-coeff", type=str, @@ -137,7 +144,7 @@ # (the regular artifact handles nswap==0). lammps.processors(args.processors) lammps.units("metal") -lammps.boundary("p p p") +lammps.boundary(args.boundary) lammps.atom_style("atomic") # ``atom_modify map yes`` is required when single-rank dispatch goes # through the regular artifact of a use_loc_mapping=False .pt2: the diff --git a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py new file mode 100644 index 0000000000..4829a91a7d --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py @@ -0,0 +1,439 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test LAMMPS with a .pt2 DPA4C model that carries a long-range module. + +The latent-charge module sums over every pair of atoms of a complete +non-periodic system, and over the images of the cell of a periodic one when +it serves periodic systems, so ``pair_style deepmd`` runs a model whose module +serves no periodic system with non-periodic boundaries only, and any +long-range model with boundaries periodic in every dimension or in none. A +compressed model also runs on several ranks, whose atoms it gathers over MPI; +an uncompressed one runs on a single rank. The models and the reference +values are generated by source/tests/infer/gen_dpa4c_long_range.py, the +compressed models for CUDA only: without a cell, at the default charge state +and at ``charge_spin 1.0 2.0``, for a module that serves no periodic system, +and in its cell and without one for a module that does. +""" + +import importlib.util +import os +import shutil +from pathlib import ( + Path, +) + +import constants +import numpy as np +import pytest +from expected_ref import ( + read_expected_ref, +) +from lammps import ( + PyLammps, +) +from lammps_test_utils import ( + make_atomic_lammps, + run_mpi_pair_runner, +) +from write_lmp_data import ( + write_lmp_data, +) + +pt2_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4c_long_range.pt2" +) +ref_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4c_long_range.expected" +) +compressed_pt2_file = pt2_file.with_name("deeppot_dpa4c_long_range_compressed.pt2") +periodic_pt2_file = pt2_file.with_name("deeppot_dpa4c_long_range_periodic.pt2") +periodic_compressed_pt2_file = pt2_file.with_name( + "deeppot_dpa4c_long_range_periodic_compressed.pt2" +) +# A model without a long-range module, for a model deviation beside one. +short_range_file = pt2_file.with_name("deeppot_sea.pth") +compressed_ref_file = ref_file.with_name("deeppot_dpa4c_long_range_compressed.expected") +data_file = Path(__file__).parent / "data_dpa4c_long_range_pt2.lmp" +periodic_data_file = Path(__file__).parent / "data_dpa4c_long_range_periodic_pt2.lmp" +replicated_data_file = Path(__file__).parent / "data_dpa4c_long_range_grid_pt2.lmp" +mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py" + +# Guarded with try/except because the generator only runs when PyTorch is +# built; matrices that disable PyTorch still load this file at collection time. +try: + _ref = read_expected_ref(ref_file) + expected_e_default = float(np.sum(_ref["default"]["expected_e"])) + expected_f_default = _ref["default"]["expected_f"].reshape(6, 3) + expected_e_explicit = float(np.sum(_ref["explicit"]["expected_e"])) + expected_f_explicit = _ref["explicit"]["expected_f"].reshape(6, 3) +except FileNotFoundError: + expected_e_default = expected_f_default = None + expected_e_explicit = expected_f_explicit = None + + +def _read_ref(path: Path) -> dict | None: + try: + return read_expected_ref(path) + except FileNotFoundError: + return None + + +compressed_ref = _read_ref(compressed_ref_file) +periodic_ref = _read_ref( + ref_file.with_name("deeppot_dpa4c_long_range_periodic.expected") +) +periodic_compressed_ref = _read_ref( + ref_file.with_name("deeppot_dpa4c_long_range_periodic_compressed.expected") +) + +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +# The triclinic cell of the periodic references: xlo xhi ylo yhi zlo zhi xy xz +# yz of the lattice vectors (13, 0, 0), (0.6, 12.5, 0) and (-0.4, 0.8, 13.5). +periodic_box = np.array([0, 13, 0, 12.5, 0, 13.5, 0.6, -0.4, 0.8]) +coord = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +# atype [0, 1, 1, 0, 1, 1] (O, H, H, O, H, H) -> LAMMPS types [1, 2, 2, 1, 2, 2] +type_OH = np.array([1, 2, 2, 1, 2, 2]) +# LAMMPS centroid/stress/atom reports the per-atom stress, the negative of the +# per-atom virial, with the components xx, yy, zz, xy, xz, yz, yx, zx, zy; these +# columns of it give the nine components of the reference virial. +expected_v_to_lammps = [0, 6, 7, 3, 1, 8, 4, 5, 2] + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip("Skip test because PyTorch support is not enabled.") + write_lmp_data(box, coord, type_OH, data_file) + write_lmp_data(periodic_box, coord, type_OH, periodic_data_file) + + +def teardown_module() -> None: + for path in (data_file, periodic_data_file, replicated_data_file): + if path.exists(): + os.remove(path) + + +def _lammps(boundary: str, data: Path = data_file) -> PyLammps: + if not pt2_file.is_file(): + pytest.skip(f"Skip because {pt2_file} was not generated.") + return make_atomic_lammps(data, boundary=boundary, atom_map="yes") + + +@pytest.fixture +def lammps(): + lmp = _lammps("f f f") + yield lmp + lmp.close() + + +@pytest.fixture +def lammps_periodic(): + lmp = _lammps("p p p") + yield lmp + lmp.close() + + +def test_pair_deepmd(lammps) -> None: + lammps.pair_style(f"deepmd {pt2_file.resolve()}") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_default) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_default[lammps.atoms[ii].id - 1] + ) + lammps.run(1) + + +def test_pair_deepmd_charge_spin(lammps) -> None: + lammps.pair_style(f"deepmd {pt2_file.resolve()} charge_spin 1.0 2.0") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_explicit) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_explicit[lammps.atoms[ii].id - 1] + ) + + +def test_periodic_boundary_is_rejected(lammps_periodic) -> None: + lammps_periodic.pair_style(f"deepmd {pt2_file.resolve()}") + lammps_periodic.pair_coeff("* *") + with pytest.raises(Exception, match="long-range module"): + lammps_periodic.run(0) + + +@pytest.mark.parametrize("long_range_first", [True, False]) +def test_model_deviation_periodic_boundary_is_rejected( + lammps_periodic, long_range_first: bool, tmp_path: Path +) -> None: + """Any model of a model deviation with a long-range module confines it.""" + if not short_range_file.is_file(): + pytest.skip(f"Skip because {short_range_file} was not generated.") + models = [pt2_file, short_range_file] + if not long_range_first: + models.reverse() + lammps_periodic.pair_style( + "deepmd " + + " ".join(str(model.resolve()) for model in models) + + f" out_freq 1 out_file {tmp_path / 'model_devi.out'}" + ) + lammps_periodic.pair_coeff("* *") + with pytest.raises(Exception, match="long-range module"): + lammps_periodic.run(0) + + +@pytest.mark.parametrize("long_range_first", [True, False]) +def test_model_deviation_without_cell_needs_long_range_throughout( + lammps, long_range_first: bool, tmp_path: Path +) -> None: + """A non-periodic model deviation passes no cell, which every model must serve.""" + if not short_range_file.is_file(): + pytest.skip(f"Skip because {short_range_file} was not generated.") + models = [pt2_file, short_range_file] + if not long_range_first: + models.reverse() + lammps.pair_style( + "deepmd " + + " ".join(str(model.resolve()) for model in models) + + f" out_freq 1 out_file {tmp_path / 'model_devi.out'}" + ) + lammps.pair_coeff("* *") + with pytest.raises(Exception, match="every model carries one"): + lammps.run(0) + + +def test_model_deviation_of_long_range_models(lammps, tmp_path: Path) -> None: + """Models that all carry a long-range module serve a non-periodic deviation.""" + md_file = tmp_path / "model_devi.out" + lammps.pair_style( + f"deepmd {pt2_file.resolve()} {pt2_file.resolve()} " + f"out_freq 1 out_file {md_file}" + ) + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_default) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_default[lammps.atoms[ii].id - 1] + ) + md = np.loadtxt(md_file) + assert md[1:] == pytest.approx(np.zeros(md.shape[0] - 1), abs=1e-10) + + +def _check_reference(lmp: PyLammps, ref: dict, section: str) -> None: + """Compare the energy and the forces of a run with a reference section.""" + lmp.run(0) + assert lmp.eval("pe") == pytest.approx(float(np.sum(ref[section]["expected_e"]))) + expected_f = ref[section]["expected_f"].reshape(6, 3) + for ii in range(6): + assert lmp.atoms[ii].force == pytest.approx( + expected_f[lmp.atoms[ii].id - 1], abs=1e-8 + ) + + +@pytest.mark.parametrize( + ("boundary", "section"), [("p p p", "periodic"), ("f f f", "cluster")] +) +def test_pair_deepmd_periodic_module(boundary: str, section: str) -> None: + """A module that serves periodic systems sums over the images of the cell.""" + if periodic_ref is None or not periodic_pt2_file.is_file(): + pytest.skip(f"Skip because {periodic_pt2_file} was not generated.") + lmp = _lammps(boundary, periodic_data_file) + lmp.pair_style(f"deepmd {periodic_pt2_file.resolve()}") + lmp.pair_coeff("* *") + _check_reference(lmp, periodic_ref, section) + lmp.run(1) + lmp.close() + + +@pytest.mark.parametrize("model", [pt2_file, periodic_pt2_file]) +def test_partially_periodic_boundary_is_rejected(model: Path) -> None: + if not model.is_file(): + pytest.skip(f"Skip because {model} was not generated.") + lmp = _lammps("p p f", periodic_data_file) + lmp.pair_style(f"deepmd {model.resolve()}") + lmp.pair_coeff("* *") + with pytest.raises(Exception, match="every dimension"): + lmp.run(0) + lmp.close() + + +mpi_available = pytest.mark.skipif( + shutil.which("mpirun") is None or importlib.util.find_spec("mpi4py") is None, + reason="MPI or mpi4py is not installed", +) + + +@mpi_available +@pytest.mark.parametrize( + ("nprocs", "processors", "section", "pair_extra"), + [ + # Both ranks own atoms. + (2, "2 1 1", "default", ""), + # Every atom lies in the first slab along z: two ranks hold ghost atoms + # alone and the last none at all, and all of them join the exchange. + (4, "1 1 4", "explicit", "charge_spin 1.0 2.0"), + ], +) +def test_pair_deepmd_mpi_compressed( + nprocs: int, processors: str, section: str, pair_extra: str +) -> None: + if compressed_ref is None or not compressed_pt2_file.is_file(): + pytest.skip(f"Skip because {compressed_pt2_file} was not generated.") + runner_args = ["--boundary", "f f f"] + if pair_extra: + runner_args += ["--pair-extra", pair_extra] + result = run_mpi_pair_runner( + mpi_runner, + data_file, + compressed_pt2_file, + nprocs=nprocs, + processors=processors, + runner_args=runner_args, + ) + expected_e = float(np.sum(compressed_ref[section]["expected_e"])) + expected_f = compressed_ref[section]["expected_f"].reshape(6, 3) + expected_v = -compressed_ref[section]["expected_v"].reshape(6, 9) + assert result["pe"] == pytest.approx(expected_e, rel=1e-6) + np.testing.assert_allclose(result["forces"], expected_f, rtol=0, atol=1e-5) + np.testing.assert_allclose( + result["virials"][:, expected_v_to_lammps] / constants.nktv2p, + expected_v, + rtol=0, + atol=1e-5, + ) + + +@mpi_available +@pytest.mark.parametrize( + ("nprocs", "processors", "boundary", "section"), + [ + (2, "2 1 1", "p p p", "periodic"), + # Every atom lies in the lower half along z; the ranks above hold + # ghost atoms alone. + (4, "1 1 4", "p p p", "periodic"), + (2, "1 2 1", "f f f", "cluster"), + ], +) +def test_pair_deepmd_mpi_periodic_compressed( + nprocs: int, processors: str, boundary: str, section: str +) -> None: + if periodic_compressed_ref is None or not periodic_compressed_pt2_file.is_file(): + pytest.skip(f"Skip because {periodic_compressed_pt2_file} was not generated.") + if not periodic_data_file.is_file(): + write_lmp_data(periodic_box, coord, type_OH, periodic_data_file) + result = run_mpi_pair_runner( + mpi_runner, + periodic_data_file, + periodic_compressed_pt2_file, + nprocs=nprocs, + processors=processors, + runner_args=["--boundary", boundary], + ) + ref = periodic_compressed_ref[section] + expected_v = -ref["expected_v"].reshape(6, 9) + assert result["pe"] == pytest.approx(float(np.sum(ref["expected_e"])), rel=1e-6) + np.testing.assert_allclose( + result["forces"], ref["expected_f"].reshape(6, 3), rtol=0, atol=1e-5 + ) + np.testing.assert_allclose( + result["virials"][:, expected_v_to_lammps] / constants.nktv2p, + expected_v, + rtol=0, + atol=1e-5, + ) + + +@pytest.fixture(scope="module") +def replicated_reference() -> dict: + """A grid of 168 copies of the six atoms, and its single-rank result. + + The copies lie farther apart than the short-range cutoff, so only the + long-range module couples them; the 1008 atoms fill several blocks and + chunks of the long-range operator. + """ + if compressed_ref is None or not compressed_pt2_file.is_file(): + pytest.skip(f"Skip because {compressed_pt2_file} was not generated.") + shifts = np.stack( + np.meshgrid( + 18.0 * np.arange(4), 6.0 * np.arange(7), 6.0 * np.arange(6), indexing="ij" + ), + axis=-1, + ).reshape(-1, 1, 3) + grid_coord = (coord[None] + shifts).reshape(-1, 3) + grid_type = np.tile(type_OH, shifts.shape[0]) + # The box reaches beyond the atoms along z, so that the last of four slabs + # holds no atom, owned or ghost. + grid_box = np.array([0, 72, 0, 42, 0, 60, 0, 0, 0]) + write_lmp_data(grid_box, grid_coord, grid_type, replicated_data_file) + return run_mpi_pair_runner( + mpi_runner, + replicated_data_file, + compressed_pt2_file, + nprocs=1, + runner_args=["--boundary", "f f f"], + ) + + +@mpi_available +@pytest.mark.parametrize( + ("nprocs", "processors"), + [ + # Uneven slabs that cut through copies of the six atoms. + (3, "3 1 1"), + # The last slab holds no atom at all. + (4, "1 1 4"), + ], +) +def test_pair_deepmd_mpi_compressed_matches_single_rank( + replicated_reference: dict, nprocs: int, processors: str +) -> None: + result = run_mpi_pair_runner( + mpi_runner, + replicated_data_file, + compressed_pt2_file, + nprocs=nprocs, + processors=processors, + runner_args=["--boundary", "f f f"], + ) + assert result["pe"] == pytest.approx(replicated_reference["pe"], rel=1e-6) + np.testing.assert_allclose( + result["forces"], replicated_reference["forces"], rtol=0, atol=1e-5 + ) + np.testing.assert_allclose( + result["virials"] / constants.nktv2p, + replicated_reference["virials"] / constants.nktv2p, + rtol=0, + atol=1e-5, + ) + + +@mpi_available +def test_pair_deepmd_mpi_uncompressed_is_rejected() -> None: + if not pt2_file.is_file(): + pytest.skip(f"Skip because {pt2_file} was not generated.") + result = run_mpi_pair_runner( + mpi_runner, + data_file, + pt2_file, + nprocs=2, + processors="2 1 1", + runner_args=["--boundary", "f f f"], + capture=True, + ) + assert result["returncode"] != 0 + assert "compact canonical lower" in result["stdout"] + result["stderr"] diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index d0566c7741..3a596618b0 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -23,7 +23,8 @@ list( cpu/neighbor_search_cpu.cc dpa4c/ops.cc dpa4c/graph_compress_cpu.cc - ${DPA4C_CPU_ISA_SRC}) + ${DPA4C_CPU_ISA_SRC} + long_range/ops.cc) # The architecture level is pinned per unit rather than added to it: the project # may be configured with ENABLE_NATIVE_OPTIMIZATION, and a trailing -march wins, # so this is what keeps the fallback units free of instructions the dispatcher @@ -240,6 +241,10 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) ${DPA4C_GRAPH_COMPRESS_KERNEL_SRC} graph_fitting.cu edge_force_virial.cu + long_range/direct_sum.cu + long_range/domain_exchange.cc + long_range/latent_charge_sum.cu + long_range/mesh_sum.cu dpa1_graph_energy_force.cu dpa4/so2_conv.cu ${DPA4_SO2_CONV_KERNEL_SRC} @@ -274,7 +279,7 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) set_property(TARGET deepmd_op_pt PROPERTY CUDA_ARCHITECTURES ${DEEPMD_PT_CUDA_ARCHITECTURES}) endif() - target_link_libraries(deepmd_op_pt PRIVATE CUDA::cublasLt) + target_link_libraries(deepmd_op_pt PRIVATE CUDA::cublasLt CUDA::cufft) # libtorch headers require at least C++17, and PyTorch 2.14+ requires C++20. # Preserve an explicitly higher CUDA standard requested by the user. set(DEEPMD_PT_CUDA_STANDARD 17) @@ -314,6 +319,11 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) if(CMAKE_SYSTEM_NAME STREQUAL "Linux") target_link_options(deepmd_op_pt PRIVATE "LINKER:--no-undefined") endif() + # The direct sum of the long range evaluates every pair with the special + # function unit; without denormal handling each of its three calls is a single + # instruction. + set_source_files_properties(long_range/direct_sum.cu + PROPERTIES COMPILE_OPTIONS "--use_fast_math") endif() if(${OP_CXX_ABI_PT} EQUAL ${OP_CXX_ABI}) target_link_libraries(deepmd_op_pt PRIVATE ${LIB_DEEPMD}) diff --git a/source/op/pt/long_range/charge_sum.h b/source/op/pt/long_range/charge_sum.h new file mode 100644 index 0000000000..04245d3287 --- /dev/null +++ b/source/op/pt/long_range/charge_sum.h @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// The two evaluations behind deepmd::latent_charge_sum (schema and formulas in +// ops.cc): the direct sum over the pairs of non-periodic frames +// (direct_sum.cu) and the smooth particle-mesh Ewald sum of periodic frames +// (mesh_sum.cu). Both take float64 positions and float32 raw charges of the +// nodes of the frames on one CUDA device, apply the constraint, and evaluate +// the owner nodes [owner_begin, owner_begin + owner_count) against every node +// of their frames. Owners other than every node lie in a single frame, whose +// nodes are then the charged nodes of every domain of ``communicator``, over +// which the per-frame sums of the owners are added. + +#pragma once + +#include + +#include +#include + +#include "domain_exchange.h" + +namespace deepmd { +namespace long_range { + +// The atomic energy (owners, 1) float64, the charge gradient (owners, C), the +// force (owners, 3), the atomic virial (owners, 9) and the virial of the +// owners of every frame (F, 9), float32. +using ChargeSumOutputs = std::tuple; + +// The direct sum. ``coefficient`` holds the float64 coefficients of the +// kernel on the device. +ChargeSumOutputs direct_sum(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int owner_begin, + int owner_count, + const torch::Tensor* communicator); + +// The mesh sum. ``box`` holds the cell of every frame, float64 on the host; +// ``coefficient`` the float64 coefficients of the kernel on the device. +ChargeSumOutputs mesh_sum(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& box, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing, + int owner_begin, + int owner_count, + const torch::Tensor* communicator); + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu new file mode 100644 index 0000000000..2e0f128d85 --- /dev/null +++ b/source/op/pt/long_range/direct_sum.cu @@ -0,0 +1,979 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA kernels of the latent-charge long range of non-periodic frames, by the +// direct sum over their pairs (schema and formulas in ops.cc, the entry in +// latent_charge_sum.cu). +// +// Every node reads every node of its frame, so the memory is linear in the +// node count while the work is quadratic. One thread owns one node and walks +// the nodes of its frame in tiles staged in shared memory; the partner of a +// thread at a given step is the same for the whole block, so each tile entry +// is a broadcast read. Each pair is evaluated by both of its nodes, which +// keeps every sum in the registers of its owner: no atomic operation is +// needed, and the result does not depend on scheduling. +// +// Positions are carried relative to the first node of their frame as a +// double-float pair: a float32 head and the float32 tail of the remainder. +// The displacement of two nodes, formed as the difference of the heads plus +// the difference of the tails, keeps the float32 precision of the +// displacement itself wherever the frame lies, while every pair evaluates in +// float32; the far field of the erf kernel forms it from the heads alone (see +// evaluate_kernel). For the erf kernel the positions are scaled by alpha, so +// that the pair loop evaluates erf(x) / x and its slope without a +// coefficient, and the reduction restores the units. +// +// A thread sums one tile in float32 and adds the tile sums with compensated +// float32 arithmetic. The partner axis is split into contiguous chunks, so +// that a launch fills the device several times over; every chunk writes its +// sums, and a second pass adds the chunks in a fixed order in float64. +// +// The charge constraint and the per-frame virial reduce over the nodes of +// every frame, in slices whose float64 sums are added in a fixed order, so +// that every output is independent of scheduling. +// +// The nodes whose sums a call evaluates, its owners, are either every node +// or, across the domains of a decomposed system, the charged nodes of one +// MPI rank, which the entry gathers with those of every rank into one frame; +// each rank evaluates its own nodes against the whole frame, so that the +// quadratic work divides among the ranks. Every rank shifts the constrained +// channel by the same sum of the gathered charges; the mean potential of the +// constraint adds the sums of the ranks in rank order, and the per-frame +// virial is the share of the owners of the rank. + +#include +#include +#include + +#include +#include +#include +#include + +#include "charge_sum.h" +#include "frame_sums.cuh" + +namespace { + +// Nodes per block, which is also the length of a tile of partners. +constexpr int kThreads = 128; +// Partners evaluated together, which share every coefficient load. +constexpr int kBatch = 4; +constexpr int kMaxGaussians = 32; +// Fills of the resident block slots of the device that one launch aims for. +constexpr int kWaves = 8; +// Shortest chunk of the partner axis, in tiles. +constexpr int kMinChunkTiles = 2; + +enum KernelKind : int { kLes = 0, kSog = 1 }; + +// Sums of one node: the potential of every channel, the force direction and +// the six independent components of the virial. +template +constexpr int sum_width() { + return kChannels + 3 + 6; +} + +// Running sum with Kahan compensation. +struct CompensatedSum { + float sum = 0.f; + float error = 0.f; + + __device__ __forceinline__ void add(float value) { + const float corrected = value - error; + const float total = sum + corrected; + error = (total - sum) - corrected; + sum = total; + } +}; + +// Writes the double-float positions relative to the first node of each +// frame, scaled by alpha for the erf kernel, the frame of every node, and the +// charges of the pairs: the raw charges of the charged nodes, shifted in the +// constrained channel, and zero for every other node. ``alpha`` is a device +// scalar, or null for a kernel read in plain units; ``frame_stat`` is null +// without a constraint. +__global__ void prepare_nodes(const double* __restrict__ coord, + const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int frame_count, + int node_count, + int channels, + int channel, + const double* __restrict__ alpha, + const double* __restrict__ frame_stat, + float4* __restrict__ head, + float4* __restrict__ tail, + int* __restrict__ frame, + float* __restrict__ pair_charge) { + const int node = blockIdx.x * blockDim.x + threadIdx.x; + if (node >= node_count) { + return; + } + const double scale = alpha == nullptr ? 1.0 : *alpha; + const int f = frame_of(frame_ptr, frame_count, node); + double3 offset = make_double3(0.0, 0.0, 0.0); + if (f >= 0) { + const int64_t first = frame_ptr[f]; + offset.x = scale * (coord[3 * node + 0] - coord[3 * first + 0]); + offset.y = scale * (coord[3 * node + 1] - coord[3 * first + 1]); + offset.z = scale * (coord[3 * node + 2] - coord[3 * first + 2]); + } + const float hx = __double2float_rn(offset.x); + const float hy = __double2float_rn(offset.y); + const float hz = __double2float_rn(offset.z); + head[node] = make_float4(hx, hy, hz, 0.f); + tail[node] = make_float4(__double2float_rn(offset.x - hx), + __double2float_rn(offset.y - hy), + __double2float_rn(offset.z - hz), 0.f); + frame[node] = f; + const bool charged = f >= 0 && node_mask[node]; + const double shift = charged && frame_stat != nullptr + ? frame_stat[f * kStatWidth + kStatShift] + : 0.0; + for (int c = 0; c < channels; ++c) { + const int64_t index = static_cast(node) * channels + c; + const double value = charged ? static_cast(charge[index]) : 0.0; + pair_charge[index] = + static_cast(c == channel ? value - shift : value); + } +} + +// K and K'(r) / r of a batch of pairs at the squared distances ``r2``. +// +// The erf kernel reads positions scaled by alpha, x = alpha r, and evaluates +// erf(x) / x and its slope (2 / sqrt(pi) exp(-x^2) - erf(x) / x) / x^2, with +// erf(x) = 1 - t P(t) exp(-x^2), t = 1 / (1 + p x), and the coefficients of +// Abramowitz and Stegun 7.1.26 (absolute error below 1.5e-7), so that one +// exponential serves the kernel and its slope. The two terms of the slope +// agree to second order in x, which amplifies the error of the approximation +// by 1 / x^2; below x^2 = kSeriesSquared the kernel and its slope therefore +// come from their Taylor series in x^2, +// +// erf(x) / x = 2 / sqrt(pi) sum_n (-1)^n x^(2n) / (n! (2n + 1)), +// slope = 2 / sqrt(pi) sum_{n >= 1} (-1)^n 2n x^(2n - 2) / +// (n! (2n + 1)), +// +// whose first six terms reach float32 precision there. The relative error of +// the kernel stays below 1e-6 for every x, that of its slope below 7e-6, the +// largest just above the threshold. A warp evaluates the series only for a +// batch with a pair below the threshold: a batch that pairs a node with +// itself, or a contact shorter than 0.4 / alpha, which for an alpha of 0.5 per +// angstrom or more is shorter than every bond but that of H2. Every other +// batch evaluates the closed form without a branch. +// +// Beyond x^2 = kFarSquared the float32 kernel equals 1 / x and its slope +// -1 / x^3 bit for bit: exp(-x^2) and erfc(x) fall below half a unit in the +// last place of the terms they would correct. A warp whose batch lies there +// entirely skips the evaluation of erf, which most pairs of a large system +// do, and forms the displacements of the batch from the heads alone. The +// omitted tails move a displacement by at most 2^-24 (|x_i| + |x_j|), with x +// the alpha-scaled offset of a node from the first node of its frame: the +// rounding of float32 positions over the extent of the frame, against a +// distance of at least sqrt(kFarSquared). For a cluster of 100,000 carbon +// atoms at alpha = 0.5 per angstrom, with offsets of up to 26 units, the +// potential and the force keep the error of the exact displacements. +// +// The sum of Gaussians evaluates sum_k a_k 2^(-c_k r^2) and its slope +// sum_k s_k 2^(-c_k r^2), s_k = -2 ln(2) c_k a_k, from the staged +// coefficients ``{a_k, -c_k, s_k}``. +constexpr float kFarSquared = 20.25f; +constexpr float kSeriesSquared = 0.16f; + +// A series in x^2 with the coefficients ``c``, in increasing powers, by +// Horner's rule. +template +__device__ __forceinline__ float even_series(const float (&c)[kTerms], + float x2) { + float sum = c[kTerms - 1]; +#pragma unroll + for (int n = kTerms - 2; n >= 0; --n) { + sum = fmaf(sum, x2, c[n]); + } + return sum; +} + +// Whether every pair of the batch of every thread of the warp lies in the far +// field of the erf kernel, judged from the head displacements, whose squared +// lengths it writes to ``r2``. The sum of Gaussians has no far field. +template +__device__ __forceinline__ bool far_batch(const float (&dx)[kBatch], + const float (&dy)[kBatch], + const float (&dz)[kBatch], + float (&r2)[kBatch]) { + if constexpr (kKernel == kSog) { + return false; + } else { + float nearest = std::numeric_limits::infinity(); +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + r2[b] = fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b])); + nearest = fminf(nearest, r2[b]); + } + return __all_sync(0xffffffffu, nearest >= kFarSquared); + } +} + +// erf(x) / x and its slope by the approximation of Abramowitz and Stegun at +// the squared distances ``r2``. +__device__ __forceinline__ void les_closed_form(const float (&r2)[kBatch], + float (&value)[kBatch], + float (&slope)[kBatch]) { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float inverse = rsqrtf(r2[b]); + const float x = r2[b] * inverse; + const float t = __fdividef(1.f, fmaf(0.3275911f, x, 1.f)); + const float poly = + t * + fmaf( + t, + fmaf(t, fmaf(t, fmaf(t, 1.061405429f, -1.453152027f), 1.421413741f), + -0.284496736f), + 0.254829592f); + const float gaussian = exp2f(-1.44269504088896341f * r2[b]); + value[b] = fmaf(-poly, gaussian, 1.f) * inverse; + slope[b] = + fmaf(1.12837916709551257f, gaussian, -value[b]) * (inverse * inverse); + } +} + +template +__device__ __forceinline__ void evaluate_kernel(const float (&r2)[kBatch], + const float4* gaussians, + int gaussian_count, + float (&value)[kBatch], + float (&slope)[kBatch]) { + if constexpr (kKernel == kLes) { + float nearest = r2[0]; +#pragma unroll + for (int b = 1; b < kBatch; ++b) { + nearest = fminf(nearest, r2[b]); + } + if (!__any_sync(0xffffffffu, nearest < kSeriesSquared)) { + les_closed_form(r2, value, slope); + return; + } + constexpr float kValueSeries[6] = {1.128379167f, -0.376126389f, + 0.1128379167f, -0.02686617065f, + 0.005223977625f, -0.0008548327023f}; + constexpr float kSlopeSeries[6] = {-0.7522527781f, 0.4513516668f, + -0.1611970239f, 0.041791821f, + -0.008548327023f, 0.001446639958f}; + // The closed form reads the threshold where the series is selected. + float clamped[kBatch]; +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + clamped[b] = fmaxf(r2[b], kSeriesSquared); + } + les_closed_form(clamped, value, slope); +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + if (r2[b] < kSeriesSquared) { + value[b] = even_series(kValueSeries, r2[b]); + slope[b] = even_series(kSlopeSeries, r2[b]); + } + } + } else { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + value[b] = 0.f; + slope[b] = 0.f; + } + for (int k = 0; k < gaussian_count; ++k) { + const float4 coefficient = gaussians[k]; +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float weight = exp2f(coefficient.y * r2[b]); + value[b] = fmaf(coefficient.x, weight, value[b]); + slope[b] = fmaf(coefficient.z, weight, slope[b]); + } + } + } +} + +// Partners a tile holds besides those of the frame of the node. +enum Exclusion : int { + // Every partner belongs to the frame and differs from the node. + kNoExclusion = 0, + // The node itself is among the partners. + kExcludeSelf = 1, + // The partners may also lie in other frames. + kExcludeOutsideFrame = 2, +}; + +// The state of the node a thread owns. +template +struct OwnNode { + int index; + int low; // first node of its frame + int high; // past the last node of its frame + float4 head; + float4 tail; + float charge[kChannels]; +}; + +// Adds the unscaled potential, force direction and virial that the partners +// of one staged tile give the node to ``sum``. +template +__device__ __forceinline__ void accumulate_tile( + const float4* __restrict__ tile_head, + const float4* __restrict__ tile_tail, + const float* __restrict__ tile_charge, + const float4* __restrict__ gaussians, + int gaussian_count, + int start, + const OwnNode& own, + float (&sum)[sum_width()]) { +#pragma unroll 4 + for (int base = 0; base < kThreads; base += kBatch) { + float dx[kBatch]; + float dy[kBatch]; + float dz[kBatch]; + float r2[kBatch]; +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float4 other_head = tile_head[base + b]; + dx[b] = own.head.x - other_head.x; + dy[b] = own.head.y - other_head.y; + dz[b] = own.head.z - other_head.z; + } + float value[kBatch]; + float slope[kBatch]; + if (far_batch(dx, dy, dz, r2)) { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float inverse = rsqrtf(r2[b]); + value[b] = inverse; + slope[b] = -inverse * (inverse * inverse); + } + } else { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float4 other_tail = tile_tail[base + b]; + dx[b] += own.tail.x - other_tail.x; + dy[b] += own.tail.y - other_tail.y; + dz[b] += own.tail.z - other_tail.z; + // Both kernels are finite at r2 = 0, which a node paired with itself + // and coincident padding nodes reach with a zero weight. + r2[b] = fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b])); + } + evaluate_kernel(r2, gaussians, gaussian_count, value, slope); + } +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const int other = start + base + b; + float weight = 1.f; + if constexpr (kExclusion == kExcludeSelf) { + weight = other != own.index ? 1.f : 0.f; + } else if constexpr (kExclusion == kExcludeOutsideFrame) { + weight = (other != own.index && + static_cast(other - own.low) < + static_cast(own.high - own.low)) + ? 1.f + : 0.f; + } + float product = 0.f; +#pragma unroll + for (int c = 0; c < kChannels; ++c) { + const float other_charge = + kExclusion == kNoExclusion + ? tile_charge[(base + b) * kChannels + c] + : weight * tile_charge[(base + b) * kChannels + c]; + sum[c] = fmaf(other_charge, value[b], sum[c]); + product = fmaf(own.charge[c], other_charge, product); + } + const float strength = product * slope[b]; + const float sx = strength * dx[b]; + const float sy = strength * dy[b]; + const float sz = strength * dz[b]; + sum[kChannels + 0] += sx; + sum[kChannels + 1] += sy; + sum[kChannels + 2] += sz; + sum[kChannels + 3] = fmaf(sx, dx[b], sum[kChannels + 3]); + sum[kChannels + 4] = fmaf(sy, dy[b], sum[kChannels + 4]); + sum[kChannels + 5] = fmaf(sz, dz[b], sum[kChannels + 5]); + sum[kChannels + 6] = fmaf(sx, dy[b], sum[kChannels + 6]); + sum[kChannels + 7] = fmaf(sx, dz[b], sum[kChannels + 7]); + sum[kChannels + 8] = fmaf(sy, dz[b], sum[kChannels + 8]); + } + } +} + +// Sums, for every owner node of one block and every partner of one chunk, the +// unscaled potential, force direction and virial of the node, and writes them +// to ``partial[chunk][component][node - owner_begin]``. The owners are the +// nodes [owner_begin, owner_begin + owner_count). +template +__global__ void __launch_bounds__(kThreads) + direct_sum_kernel(const float4* __restrict__ head, + const float4* __restrict__ tail, + const float* __restrict__ charge, + const int* __restrict__ frame, + const int64_t* __restrict__ frame_ptr, + int owner_begin, + int owner_count, + const double* __restrict__ coefficient, + int gaussian_count, + float* __restrict__ partial) { + constexpr int kWidth = sum_width(); + __shared__ float4 tile_head[kThreads]; + __shared__ float4 tile_tail[kThreads]; + __shared__ float tile_charge[kThreads * kChannels]; + __shared__ float4 gaussians[kMaxGaussians]; + __shared__ int range[2]; + __shared__ bool single_frame; + + const int owner_end = owner_begin + owner_count; + const int first = owner_begin + blockIdx.x * kThreads; + const int node = first + threadIdx.x; + const int own_frame = node < owner_end ? frame[node] : -1; + + if constexpr (kKernel == kSog) { + for (int k = threadIdx.x; k < gaussian_count; k += kThreads) { + const float amplitude = static_cast(coefficient[k]); + const float rate = static_cast(coefficient[gaussian_count + k]); + gaussians[k] = make_float4(amplitude, -rate, + -1.38629436111989061f * rate * amplitude, 0.f); + } + } + // The partners of the block are the nodes of the frames its nodes belong + // to: a contiguous range, since frames are contiguous on the node axis and + // padding nodes trail them. + if (threadIdx.x == 0) { + int last = min(first + kThreads, owner_end) - 1; + while (last > first && frame[last] < 0) { + --last; + } + const int low_frame = frame[first]; + const int high_frame = frame[last]; + range[0] = low_frame < 0 ? 0 : static_cast(frame_ptr[low_frame]); + range[1] = low_frame < 0 ? 0 : static_cast(frame_ptr[high_frame + 1]); + // A block whose nodes all lie in one frame finds every partner of its + // range in that frame, which spares the frame test of every pair. + single_frame = low_frame >= 0 && low_frame == high_frame && + min(first + kThreads, owner_end) - 1 == last; + } + __syncthreads(); + + // The chunk of the partner range this block sums, in whole tiles. + const int64_t tile_count = (range[1] - range[0] + kThreads - 1) / kThreads; + const int begin = + range[0] + + static_cast(tile_count * blockIdx.y / gridDim.y) * kThreads; + const int end = min( + range[1], + range[0] + static_cast(tile_count * (blockIdx.y + 1) / gridDim.y) * + kThreads); + + OwnNode own; + own.index = node; + own.low = 0; + own.high = 0; + own.head = make_float4(0.f, 0.f, 0.f, 0.f); + own.tail = own.head; + if (own_frame >= 0) { + own.head = head[node]; + own.tail = tail[node]; + own.low = static_cast(frame_ptr[own_frame]); + own.high = static_cast(frame_ptr[own_frame + 1]); + } +#pragma unroll + for (int c = 0; c < kChannels; ++c) { + own.charge[c] = own_frame >= 0 ? charge[node * kChannels + c] : 0.f; + } + + CompensatedSum total[kWidth]; + for (int start = begin; start < end; start += kThreads) { + const int partner = start + threadIdx.x; + const bool staged = partner < end; + // A slot past the chunk carries a zero charge and contributes nothing. + tile_head[threadIdx.x] = + staged ? head[partner] : make_float4(0.f, 0.f, 0.f, 0.f); + tile_tail[threadIdx.x] = + staged ? tail[partner] : make_float4(0.f, 0.f, 0.f, 0.f); +#pragma unroll + for (int c = 0; c < kChannels; ++c) { + tile_charge[threadIdx.x * kChannels + c] = + staged ? charge[partner * kChannels + c] : 0.f; + } + __syncthreads(); + + float sum[kWidth]; +#pragma unroll + for (int k = 0; k < kWidth; ++k) { + sum[k] = 0.f; + } + // In a single-frame block only the tiles that overlap the block's own + // nodes hold a node itself, so the tests are confined to them and to the + // blocks that span frames; the choice is uniform over the block. + if (!single_frame) { + accumulate_tile( + tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, + own, sum); + } else if (start < first + kThreads && first < start + kThreads) { + accumulate_tile( + tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, + own, sum); + } else { + accumulate_tile( + tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, + own, sum); + } +#pragma unroll + for (int k = 0; k < kWidth; ++k) { + total[k].add(sum[k]); + } + __syncthreads(); + } + + if (node < owner_end) { + float* out = partial + + static_cast(blockIdx.y) * kWidth * owner_count + + (node - owner_begin); +#pragma unroll + for (int k = 0; k < kWidth; ++k) { + out[static_cast(k) * owner_count] = total[k].sum; + } + } +} + +// Adds the chunk sums of one component of one owner node in chunk order, one +// thread per component and owner, and scales the result into the potential, +// the force or the per-node virial of the owner. With positions scaled by +// alpha, K(r) = alpha f(x), K'(r) / r = alpha^3 f'(x) / x and d = d' / alpha, +// so the potential carries alpha, the force alpha^2 and the virial alpha. +template +__global__ void reduce_chunks(const float* __restrict__ partial, + int owner_count, + int chunk_count, + double prefactor, + const double* __restrict__ alpha, + float* __restrict__ potential, + float* __restrict__ force, + float* __restrict__ atom_virial) { + constexpr int kWidth = sum_width(); + const int64_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= static_cast(kWidth) * owner_count) { + return; + } + const int component = static_cast(index / owner_count); + const int node = static_cast(index - component * int64_t(owner_count)); + double value = 0.0; + for (int chunk = 0; chunk < chunk_count; ++chunk) { + value += + partial[static_cast(chunk) * kWidth * owner_count + index]; + } + const double unit = alpha == nullptr ? 1.0 : *alpha; + if (component < kChannels) { + potential[node * kChannels + component] = + static_cast(prefactor * unit * value); + } else if (component < kChannels + 3) { + force[node * 3 + component - kChannels] = + static_cast(-prefactor * unit * unit * value); + } else { + // Each pair term is split half to each of its two nodes. The six sums + // are xx, yy, zz, xy, xz, yz. + constexpr int kRow[6] = {0, 1, 2, 0, 0, 1}; + constexpr int kColumn[6] = {0, 1, 2, 1, 2, 2}; + const int pair = component - kChannels - 3; + const float entry = static_cast(-0.5 * prefactor * unit * value); + float* out = atom_virial + static_cast(node) * 9; + out[3 * kRow[pair] + kColumn[pair]] = entry; + out[3 * kColumn[pair] + kRow[pair]] = entry; + } +} + +// Sums, over one slice of one frame, the per-node virial and, with a +// constraint, the potential of the charged nodes in the constrained channel. +template +__global__ void __launch_bounds__(kFrameThreads) + virial_slices(const float* __restrict__ atom_virial, + const float* __restrict__ potential, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int channels, + int channel, + double* __restrict__ slice_sum) { + constexpr int kComponents = kConstrained ? 10 : 9; + const NodeRange slice = slice_of_frame(frame_ptr); + double value[kComponents] = {}; + for (int64_t node = slice.begin + threadIdx.x; node < slice.end; + node += kFrameThreads) { +#pragma unroll + for (int k = 0; k < 9; ++k) { + value[k] += atom_virial[node * 9 + k]; + } + if constexpr (kConstrained) { + if (node_mask[node]) { + value[9] += potential[node * channels + channel]; + } + } + } + write_slice_sum(value, slice_sum); +} + +// Adds the slice sums of virial_slices into the virial of every frame and, +// with a constraint, the potential summed over its charged nodes. +template +__global__ void virial_total(const double* __restrict__ slice_sum, + int slice_count, + float* __restrict__ virial, + double* __restrict__ frame_stat) { + constexpr int kComponents = kConstrained ? 10 : 9; + if (threadIdx.x >= kComponents) { + return; + } + double value = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + value += + slice_sum[(static_cast(blockIdx.x) * slice_count + slice) * + kComponents + + threadIdx.x]; + } + if (threadIdx.x < 9) { + virial[blockIdx.x * 9 + threadIdx.x] = static_cast(value); + } else { + frame_stat[blockIdx.x * kStatWidth + kStatPotential] = value; + } +} + +// Writes the atomic energy q_i . phi_i / 2 of every owner node and the +// derivative of the energy with respect to its raw charges: w_i (phi_i - +// mean_f) in the constrained channel and w_i phi_i in the others, w_i marking +// the charged nodes. ``pair_charge``, ``node_mask`` and ``frame`` start at the +// first owner. +__global__ void project_charges(const float* __restrict__ potential, + const float* __restrict__ pair_charge, + const bool* __restrict__ node_mask, + const int* __restrict__ frame, + const double* __restrict__ frame_stat, + int owner_count, + int channels, + int channel, + double* __restrict__ atom_energy, + float* __restrict__ charge_gradient) { + const int node = blockIdx.x * blockDim.x + threadIdx.x; + if (node >= owner_count) { + return; + } + const int f = frame[node]; + const bool charged = f >= 0 && node_mask[node]; + const double mean = charged && frame_stat != nullptr + ? frame_stat[f * kStatWidth + kStatPotential] / + frame_stat[f * kStatWidth + kStatCount] + : 0.0; + double energy = 0.0; + for (int c = 0; c < channels; ++c) { + const int64_t index = static_cast(node) * channels + c; + const double phi = potential[index]; + energy += static_cast(pair_charge[index]) * phi; + const double gradient = charged ? phi : 0.0; + charge_gradient[index] = + static_cast(c == channel ? gradient - mean : gradient); + } + atom_energy[node] = 0.5 * energy; +} + +// Chunks of the partner axis for the given counts of owners and partners. The +// blocks of one launch all carry the same work, so the launch fills the +// resident block slots of the device several times over, which bounds the +// idle share of the last wave; no chunk is shorter than kMinChunkTiles tiles +// of a frame that spans every partner. +int chunk_count_for(int owner_count, + int partner_count, + int resident_blocks_per_multiprocessor) { + const int multiprocessors = + at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + const int owner_blocks = (owner_count + kThreads - 1) / kThreads; + const int partner_tiles = (partner_count + kThreads - 1) / kThreads; + const int64_t wanted_blocks = static_cast(kWaves) * + resident_blocks_per_multiprocessor * + multiprocessors; + const int wanted = + static_cast((wanted_blocks + owner_blocks - 1) / owner_blocks); + return std::max(1, std::min(wanted, partner_tiles / kMinChunkTiles)); +} + +// The staged partners, one row per node, and the sums of the owners. +struct DirectSumBuffers { + const float4* head; + const float4* tail; + const float* charge; + const int* frame; + const int64_t* frame_ptr; + int node_count; + const double* coefficient; + int gaussian_count; + float* potential; + float* force; + float* atom_virial; +}; + +template +int resident_blocks_per_multiprocessor() { + // The occupancy is a property of the compiled kernel and of the device; a + // process drives one device architecture. + static const int blocks = [] { + int value = 0; + C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &value, direct_sum_kernel, kThreads, 0)); + return std::max(1, value); + }(); + return blocks; +} + +template +void launch(const DirectSumBuffers& buffers, + int owner_begin, + int owner_count, + double prefactor, + cudaStream_t stream) { + constexpr int kWidth = sum_width(); + const int chunk_count = + chunk_count_for(owner_count, buffers.node_count, + resident_blocks_per_multiprocessor()); + auto partial = + torch::empty({static_cast(chunk_count) * owner_count * kWidth}, + torch::TensorOptions() + .dtype(torch::kFloat32) + .device(torch::kCUDA, c10::cuda::current_device())); + const dim3 grid((owner_count + kThreads - 1) / kThreads, chunk_count); + direct_sum_kernel<<>>( + buffers.head, buffers.tail, buffers.charge, buffers.frame, + buffers.frame_ptr, owner_begin, owner_count, buffers.coefficient, + buffers.gaussian_count, partial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const int64_t reduce_threads = static_cast(kWidth) * owner_count; + reduce_chunks + <<((reduce_threads + 255) / 256), 256, 0, stream>>>( + partial.data_ptr(), owner_count, chunk_count, prefactor, + kKernel == kLes ? buffers.coefficient : nullptr, buffers.potential, + buffers.force, buffers.atom_virial); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void dispatch_channels(int channels, + const DirectSumBuffers& buffers, + int owner_begin, + int owner_count, + double prefactor, + cudaStream_t stream) { + switch (channels) { + case 1: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 2: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 3: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 4: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 5: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 6: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 7: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 8: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + default: + TORCH_CHECK(false, "latent_charge_sum: unsupported channel count ", + channels); + } +} + +} // namespace + +namespace deepmd { +namespace long_range { + +// Evaluates the owner nodes [owner_begin, owner_begin + owner_count) against +// the nodes of their frames. Returns the outputs of the owners, row r for +// owner owner_begin + r, and the per-frame virial of the owners. Owners other +// than every node lie in a single frame. A non-null ``communicator`` spans +// the domains whose owners share that frame, over which the potential of the +// constraint is summed. +ChargeSumOutputs direct_sum(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int owner_begin, + int owner_count, + const torch::Tensor* communicator) { + TORCH_CHECK(coord.size(0) < std::numeric_limits::max() / kThreads, + "latent_charge_sum: too many nodes"); + const int node_count = static_cast(coord.size(0)); + const int frame_count = static_cast(n_node_per_frame.size(0)); + const int64_t channels = charge.size(1); + const auto float_options = charge.options(); + const auto double_options = coord.options(); + if (node_count == 0 || frame_count == 0) { + return {torch::zeros({owner_count, 1}, double_options), + torch::zeros({owner_count, channels}, float_options), + torch::zeros({owner_count, 3}, float_options), + torch::zeros({owner_count, 9}, float_options), + torch::zeros({frame_count, 9}, float_options)}; + } + const bool every_node = owner_begin == 0 && owner_count == node_count; + TORCH_CHECK(every_node || frame_count == 1, + "latent_charge_sum: owners other than every node lie in " + "one frame"); + const bool constrained = channel >= 0; + const int constrained_channel = static_cast(channel); + // The kernels write every entry of the outputs. + auto atom_energy = torch::empty({owner_count, 1}, double_options); + auto charge_gradient = torch::empty({owner_count, channels}, float_options); + auto force = torch::empty({owner_count, 3}, float_options); + auto atom_virial = torch::empty({owner_count, 9}, float_options); + auto virial = torch::empty({frame_count, 9}, float_options); + + const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); + const auto slices_for = [frame_count](int64_t nodes) { + const int64_t average_nodes = (nodes + frame_count - 1) / frame_count; + return static_cast(std::min( + kMaxSlices, + std::max(1, (average_nodes + kSliceNodes - 1) / kSliceNodes))); + }; + const int slice_count = slices_for(node_count); + const int owner_slice_count = slices_for(owner_count); + const int components = constrained ? 10 : 9; + Workspace workspace; + const int64_t frame_ptr_at = workspace.reserve(frame_count + 1); + const int64_t head_at = workspace.reserve(node_count); + const int64_t tail_at = workspace.reserve(node_count); + const int64_t frame_at = workspace.reserve(node_count); + const int64_t pair_charge_at = + workspace.reserve(int64_t{node_count} * channels); + const int64_t potential_at = + workspace.reserve(int64_t{owner_count} * channels); + const int64_t frame_stat_at = + workspace.reserve(int64_t{frame_count} * kStatWidth); + const int64_t slice_charge_at = + workspace.reserve(int64_t{frame_count} * slice_count * 2); + const int64_t slice_virial_at = workspace.reserve( + int64_t{frame_count} * owner_slice_count * components); + workspace.allocate(coord.device()); + auto* frame_ptr = workspace.at(frame_ptr_at); + auto* head = workspace.at(head_at); + auto* tail = workspace.at(tail_at); + auto* frame = workspace.at(frame_at); + auto* pair_charge = workspace.at(pair_charge_at); + auto* potential = workspace.at(potential_at); + auto* frame_stat = + constrained ? workspace.at(frame_stat_at) : nullptr; + auto* slice_charge = workspace.at(slice_charge_at); + auto* slice_virial = workspace.at(slice_virial_at); + + const auto stream = at::cuda::getCurrentCUDAStream(); + frame_offsets<<<1, kOffsetThreads, 0, stream>>>(counts.data_ptr(), + frame_count, frame_ptr); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (constrained) { + charge_slices<<>>( + charge.data_ptr(), node_mask.data_ptr(), frame_ptr, + static_cast(channels), constrained_channel, slice_charge); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads, + kFrameThreads, 0, stream>>>( + slice_charge, slice_count, frame_count, total_charge.data_ptr(), + frame_stat); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // Alpha stays on the device, so that a call never waits for it. + const double* alpha = + kernel == kLes ? coefficient.data_ptr() : nullptr; + prepare_nodes<<<(node_count + kThreads - 1) / kThreads, kThreads, 0, + stream>>>(coord.data_ptr(), charge.data_ptr(), + node_mask.data_ptr(), frame_ptr, frame_count, + node_count, static_cast(channels), + constrained_channel, alpha, frame_stat, head, tail, + frame, pair_charge); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + if (owner_count > 0) { + const DirectSumBuffers buffers{head, + tail, + pair_charge, + frame, + frame_ptr, + node_count, + coefficient.data_ptr(), + static_cast(coefficient.numel() / 2), + potential, + force.data_ptr(), + atom_virial.data_ptr()}; + if (kernel == kLes) { + dispatch_channels(static_cast(channels), buffers, owner_begin, + owner_count, prefactor, stream.stream()); + } else { + dispatch_channels(static_cast(channels), buffers, owner_begin, + owner_count, prefactor, stream.stream()); + } + } + // The frames of the owners, on the axis of the owners: the frames + // themselves when every node is an owner, else the one frame the owners + // share, written on the device. + torch::Tensor owner_frame_tensor; + const int64_t* owner_frame_ptr = frame_ptr; + if (!every_node) { + owner_frame_tensor = + torch::full({2}, int64_t{owner_count}, counts.options()); + owner_frame_tensor.select(0, 0).zero_(); + owner_frame_ptr = owner_frame_tensor.data_ptr(); + } + const bool* owner_mask = node_mask.data_ptr() + owner_begin; + const dim3 owner_slice_grid(frame_count, owner_slice_count); + if (constrained) { + virial_slices<<>>( + atom_virial.data_ptr(), potential, owner_mask, owner_frame_ptr, + static_cast(channels), constrained_channel, slice_virial); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + virial_total<<>>( + slice_virial, owner_slice_count, virial.data_ptr(), frame_stat); + } else { + virial_slices<<>>( + atom_virial.data_ptr(), potential, owner_mask, owner_frame_ptr, + static_cast(channels), constrained_channel, slice_virial); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + virial_total<<>>( + slice_virial, owner_slice_count, virial.data_ptr(), nullptr); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (constrained && communicator != nullptr) { + // The mean potential of the constraint runs over the charged nodes of + // every domain. + auto stat = + torch::from_blob(frame_stat, {frame_count, kStatWidth}, double_options); + auto potential_sum = stat.select(1, kStatPotential); + potential_sum.copy_( + sum_over_domains(*communicator, potential_sum.contiguous())); + } + if (owner_count > 0) { + project_charges<<<(owner_count + kThreads - 1) / kThreads, kThreads, 0, + stream>>>( + potential, pair_charge + int64_t{owner_begin} * channels, owner_mask, + frame + owner_begin, frame_stat, owner_count, + static_cast(channels), constrained_channel, + atom_energy.data_ptr(), charge_gradient.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + return {atom_energy, charge_gradient, force, atom_virial, virial}; +} + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/domain_exchange.cc b/source/op/pt/long_range/domain_exchange.cc new file mode 100644 index 0000000000..2ae3d06e05 --- /dev/null +++ b/source/op/pt/long_range/domain_exchange.cc @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// MPI exchanges of the latent-charge long range (contract in +// domain_exchange.h). +// +// With an MPI library that reads device memory (MPIX_Query_cuda_support), the +// rows travel between device buffers once the current stream has finished +// writing them; otherwise they are staged through pinned host buffers. Either +// way the host waits for the device once per exchange, before MPI reads the +// data; the results return to the device without a further wait. + +#include "domain_exchange.h" + +#ifdef USE_MPI +#include +#ifdef OMPI_MPI_H +#include +#endif +#endif +#include +#include + +#include +#include + +namespace deepmd { +namespace long_range { +namespace { + +std::int64_t address_of(const torch::Tensor& communicator) { + TORCH_CHECK(communicator.device().is_cpu() && + communicator.scalar_type() == torch::kInt64 && + communicator.numel() == 1, + "latent_charge_sum: communicator must be an int64 CPU " + "tensor with one element"); + return communicator.contiguous().data_ptr()[0]; +} + +#ifdef USE_MPI +// The handle of the caller, of the width of MPI_Comm in this MPI library: an +// int in MPICH, a pointer in Open MPI. +MPI_Comm communicator_of(const torch::Tensor& communicator) { + static_assert(sizeof(std::uintptr_t) <= sizeof(std::int64_t), + "an address must fit in an int64 tensor"); + return *reinterpret_cast( + static_cast(address_of(communicator))); +} + +bool reads_device_memory() { +#ifdef NO_CUDA_AWARE + return false; +#else + return MPIX_Query_cuda_support() != 0; +#endif +} +#endif + +} // namespace + +Domain domain_of(const torch::Tensor& communicator) { + if (address_of(communicator) == 0) { + return {0, 1}; + } +#ifdef USE_MPI + int initialized = 0; + MPI_Initialized(&initialized); + TORCH_CHECK(initialized != 0, + "latent_charge_sum: a communicator was given, but MPI is " + "not initialized"); + const MPI_Comm comm = communicator_of(communicator); + Domain domain{0, 1}; + MPI_Comm_rank(comm, &domain.rank); + MPI_Comm_size(comm, &domain.size); + return domain; +#else + TORCH_CHECK(false, + "latent_charge_sum: an evaluation across MPI ranks " + "requires DeePMD-kit built with MPI"); +#endif +} + +std::pair gather_domain_rows( + const torch::Tensor& communicator, const torch::Tensor& rows) { +#ifdef USE_MPI + TORCH_CHECK(rows.is_cuda() && rows.dim() == 2 && + rows.scalar_type() == torch::kFloat64 && rows.is_contiguous(), + "gather_domain_rows: rows must be a contiguous float64 CUDA " + "matrix"); + const MPI_Comm comm = communicator_of(communicator); + int rank = 0; + int size = 1; + MPI_Comm_rank(comm, &rank); + MPI_Comm_size(comm, &size); + const std::int64_t width = rows.size(1); + std::int64_t local = rows.size(0); + std::vector counts(size); + MPI_Allgather(&local, 1, MPI_INT64_T, counts.data(), 1, MPI_INT64_T, comm); + std::vector receive_counts(size); + std::vector displacements(size); + std::int64_t total = 0; + std::int64_t offset = 0; + for (int r = 0; r < size; ++r) { + if (r == rank) { + offset = total; + } + TORCH_CHECK((total + counts[r]) * width <= std::numeric_limits::max(), + "gather_domain_rows: the rows of all ranks exceed the int " + "range of an MPI count"); + receive_counts[r] = static_cast(counts[r] * width); + displacements[r] = static_cast(total * width); + total += counts[r]; + } + const int send_count = static_cast(local * width); + auto gathered = torch::empty({total, width}, rows.options()); + if (reads_device_memory()) { + at::cuda::getCurrentCUDAStream().synchronize(); + MPI_Allgatherv(rows.data_ptr(), send_count, MPI_DOUBLE, + gathered.data_ptr(), receive_counts.data(), + displacements.data(), MPI_DOUBLE, comm); + } else { + const auto host = rows.options().device(torch::kCPU).pinned_memory(true); + auto host_rows = torch::empty({local, width}, host); + host_rows.copy_(rows); + auto host_gathered = torch::empty({total, width}, host); + MPI_Allgatherv(host_rows.data_ptr(), send_count, MPI_DOUBLE, + host_gathered.data_ptr(), receive_counts.data(), + displacements.data(), MPI_DOUBLE, comm); + gathered.copy_(host_gathered, /*non_blocking=*/true); + } + return {gathered, offset}; +#else + TORCH_CHECK(false, + "gather_domain_rows: an evaluation across MPI ranks requires " + "DeePMD-kit built with MPI"); +#endif +} + +torch::Tensor sum_over_domains(const torch::Tensor& communicator, + const torch::Tensor& values) { +#ifdef USE_MPI + TORCH_CHECK(values.dim() == 1 && values.scalar_type() == torch::kFloat64, + "sum_over_domains: values must be a float64 vector"); + const MPI_Comm comm = communicator_of(communicator); + int size = 1; + MPI_Comm_size(comm, &size); + const torch::Tensor local = values.to(torch::kCPU).contiguous(); + const int count = static_cast(local.numel()); + std::vector every(static_cast(size) * count); + MPI_Allgather(local.data_ptr(), count, MPI_DOUBLE, every.data(), + count, MPI_DOUBLE, comm); + auto total = torch::zeros({count}, local.options().pinned_memory(true)); + double* sum = total.data_ptr(); + for (int r = 0; r < size; ++r) { + for (int k = 0; k < count; ++k) { + sum[k] += every[static_cast(r) * count + k]; + } + } + return total.to(values.device(), /*non_blocking=*/true); +#else + TORCH_CHECK(false, + "sum_over_domains: an evaluation across MPI ranks requires " + "DeePMD-kit built with MPI"); +#endif +} + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/domain_exchange.h b/source/op/pt/long_range/domain_exchange.h new file mode 100644 index 0000000000..f423974f26 --- /dev/null +++ b/source/op/pt/long_range/domain_exchange.h @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Exchange of the latent-charge long range between the domains of a +// decomposed system. +// +// A domain is one MPI rank of the communicator that a caller hands to +// deepmd::latent_charge_sum as an int64 tensor holding the address of its +// MPI_Comm handle. A zero address stands for a single domain. Both exchanges +// follow the rank order, so every rank receives the same rows and the same +// sums, independent of the reduction algorithm of the MPI library. + +#pragma once + +#include + +#include +#include + +namespace deepmd { +namespace long_range { + +struct Domain { + int rank; + int size; +}; + +// The rank of this process and the number of ranks of ``communicator``; one +// rank for a zero address. +Domain domain_of(const torch::Tensor& communicator); + +// Concatenates, in rank order, the rows of every rank. ``rows`` is a +// contiguous float64 matrix of shape (M_r, W) on the CUDA device, with the +// same W on every rank. Returns the rows of all ranks, of shape (sum_r M_r, +// W) on the device of ``rows``, and the index of the first row of this rank. +std::pair gather_domain_rows( + const torch::Tensor& communicator, const torch::Tensor& rows); + +// Sums the float64 vector ``values`` over the ranks, adding the vectors of +// the ranks in rank order. Returns the sum on the device of ``values``. +torch::Tensor sum_over_domains(const torch::Tensor& communicator, + const torch::Tensor& values); + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/frame_sums.cuh b/source/op/pt/long_range/frame_sums.cuh new file mode 100644 index 0000000000..d3f8d40af6 --- /dev/null +++ b/source/op/pt/long_range/frame_sums.cuh @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Per-frame bookkeeping shared by the CUDA units of the latent-charge long +// range: the workspace of a call, the offsets of the frames on the node axis, +// the sums over the nodes of every frame in a fixed order, and the statistics +// of the charge constraint. +// +// The kernels are defined in an unnamed namespace, so that every unit that +// includes this header compiles its own copy. + +#pragma once + +#include +#include + +#include +#include + +namespace { + +// The intermediate buffers of one call, carved out of a single device +// allocation: ``reserve`` lays out a buffer of ``count`` elements and returns +// its byte offset, and ``at`` its address once ``allocate`` has made the +// allocation. Every buffer starts on a 256-byte boundary. +class Workspace { + public: + template + int64_t reserve(int64_t count) { + bytes_ = (bytes_ + kAlignment - 1) / kAlignment * kAlignment; + const int64_t offset = bytes_; + bytes_ += count * static_cast(sizeof(T)); + return offset; + } + + void allocate(const torch::Device& device) { + storage_ = torch::empty( + {std::max(bytes_, 1)}, + torch::TensorOptions().dtype(torch::kUInt8).device(device)); + } + + template + T* at(int64_t offset) const { + return reinterpret_cast(storage_.data_ptr() + offset); + } + + private: + static constexpr int64_t kAlignment = 256; + int64_t bytes_ = 0; + torch::Tensor storage_; +}; + +// Index of the frame holding ``node``, or -1 for a padding node past every +// frame. ``frame_ptr`` holds ``frame_count + 1`` offsets. +__device__ __forceinline__ int frame_of(const int64_t* frame_ptr, + int frame_count, + int64_t node) { + if (node >= frame_ptr[frame_count]) { + return -1; + } + int low = 0; + int high = frame_count; + while (high - low > 1) { + const int middle = (low + high) / 2; + if (frame_ptr[middle] <= node) { + low = middle; + } else { + high = middle; + } + } + return low; +} + +// Writes the offsets of the frames on the node axis: frame_ptr[f] counts the +// nodes of the frames before f, for f in [0, frame_count]. One block scans +// the counts, each thread a contiguous run of frames. +constexpr int kOffsetThreads = 256; + +__global__ void __launch_bounds__(kOffsetThreads) + frame_offsets(const int64_t* __restrict__ n_node_per_frame, + int frame_count, + int64_t* __restrict__ frame_ptr) { + __shared__ int64_t run_total[kOffsetThreads]; + const int run = (frame_count + kOffsetThreads - 1) / kOffsetThreads; + const int begin = min(frame_count, static_cast(threadIdx.x) * run); + const int end = min(frame_count, begin + run); + int64_t total = 0; + for (int f = begin; f < end; ++f) { + total += n_node_per_frame[f]; + } + run_total[threadIdx.x] = total; + __syncthreads(); + // Inclusive scan of the run totals. + for (int stride = 1; stride < kOffsetThreads; stride *= 2) { + const int64_t addend = + threadIdx.x >= stride ? run_total[threadIdx.x - stride] : 0; + __syncthreads(); + run_total[threadIdx.x] += addend; + __syncthreads(); + } + int64_t offset = threadIdx.x == 0 ? 0 : run_total[threadIdx.x - 1]; + for (int f = begin; f < end; ++f) { + frame_ptr[f] = offset; + offset += n_node_per_frame[f]; + } + if (threadIdx.x == kOffsetThreads - 1) { + frame_ptr[frame_count] = run_total[kOffsetThreads - 1]; + } +} + +// Per-frame sums. The nodes of a frame are split into ``gridDim.y`` +// contiguous slices, one block of kFrameThreads each, and a finalizing kernel +// adds the float64 sums of the slices in slice order. +constexpr int kFrameThreads = 256; +// Nodes per slice of a frame that a slice pass aims for. +constexpr int kSliceNodes = 2048; +constexpr int kMaxSlices = 256; + +struct NodeRange { + int64_t begin; + int64_t end; +}; + +// The nodes of the slice of the frame that this block sums. +__device__ __forceinline__ NodeRange +slice_of_frame(const int64_t* __restrict__ frame_ptr) { + const int64_t frame_begin = frame_ptr[blockIdx.x]; + const int64_t frame_end = frame_ptr[blockIdx.x + 1]; + const int64_t length = (frame_end - frame_begin + gridDim.y - 1) / gridDim.y; + const int64_t begin = frame_begin + blockIdx.y * length; + return {begin, min(frame_end, begin + length)}; +} + +// Adds ``value`` over the threads of the block in a fixed order and writes the +// totals to the entry of this slice in ``slice_sum``. +template +__device__ __forceinline__ void write_slice_sum( + const double (&value)[kComponents], double* __restrict__ slice_sum) { + __shared__ double partial[kComponents][kFrameThreads]; +#pragma unroll + for (int k = 0; k < kComponents; ++k) { + partial[k][threadIdx.x] = value[k]; + } + __syncthreads(); + for (int stride = kFrameThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { +#pragma unroll + for (int k = 0; k < kComponents; ++k) { + partial[k][threadIdx.x] += partial[k][threadIdx.x + stride]; + } + } + __syncthreads(); + } + if (threadIdx.x < kComponents) { + slice_sum[(static_cast(blockIdx.x) * gridDim.y + blockIdx.y) * + kComponents + + threadIdx.x] = partial[threadIdx.x][0]; + } +} + +// Per-frame statistics of the charge constraint: the shift of the +// constrained channel, the number of charged nodes, and the potential of the +// charged nodes in that channel summed over them. +constexpr int kStatShift = 0; +constexpr int kStatCount = 1; +constexpr int kStatPotential = 2; +constexpr int kStatWidth = 3; + +// Sums, over one slice of one frame, the number of charged nodes and their raw +// charges in the constrained channel. +__global__ void __launch_bounds__(kFrameThreads) + charge_slices(const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int channels, + int channel, + double* __restrict__ slice_sum) { + const NodeRange slice = slice_of_frame(frame_ptr); + double value[2] = {0.0, 0.0}; + for (int64_t node = slice.begin + threadIdx.x; node < slice.end; + node += kFrameThreads) { + if (node_mask[node]) { + value[0] += 1.0; + value[1] += charge[node * channels + channel]; + } + } + write_slice_sum<2>(value, slice_sum); +} + +// The shift s_f = (sum_f c - Q_f) / max(n_f, 1) of the constrained channel of +// every frame, from the slice sums of charge_slices. +__global__ void charge_shift(const double* __restrict__ slice_sum, + int slice_count, + int frame_count, + const double* __restrict__ total_charge, + double* __restrict__ frame_stat) { + const int f = blockIdx.x * blockDim.x + threadIdx.x; + if (f >= frame_count) { + return; + } + double count = 0.0; + double sum = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + const double* entry = + slice_sum + (static_cast(f) * slice_count + slice) * 2; + count += entry[0]; + sum += entry[1]; + } + // A frame without a charged node has nothing to shift. + count = fmax(count, 1.0); + frame_stat[f * kStatWidth + kStatShift] = (sum - total_charge[f]) / count; + frame_stat[f * kStatWidth + kStatCount] = count; +} + +} // namespace diff --git a/source/op/pt/long_range/latent_charge_sum.cu b/source/op/pt/long_range/latent_charge_sum.cu new file mode 100644 index 0000000000..4763c72fcc --- /dev/null +++ b/source/op/pt/long_range/latent_charge_sum.cu @@ -0,0 +1,235 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA entry of deepmd::latent_charge_sum (schema and formulas in ops.cc). +// +// A call evaluates its frames by the direct sum when they carry no cell and +// by the mesh sum when they do (charge_sum.h). Across the domains of a +// decomposed system, the ranks gather the positions and charges of the charged +// nodes of every domain into one frame, in rank order (domain_exchange.h), and +// each rank evaluates its own charged nodes against that frame; the outputs +// return to the node axis of the rank, where the nodes it does not own carry +// no charge. + +#include +#include +#include + +#include +#include +#include + +#include "charge_sum.h" +#include "domain_exchange.h" + +namespace { + +using deepmd::long_range::ChargeSumOutputs; + +constexpr int kMaxChannels = 8; +constexpr int kMaxGaussians = 32; + +enum KernelKind : int { kLes = 0, kSog = 1 }; + +// The evaluation of the frames, by the direct or the mesh sum, of the owner +// nodes [owner_begin, owner_begin + owner_count). ``box`` is defined for +// periodic frames only. +struct Evaluation { + const torch::Tensor& box; + const torch::Tensor& coefficient; + int64_t kernel; + int64_t channel; + double prefactor; + int64_t order; + double spacing; + bool periodic; + + ChargeSumOutputs operator()(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + int owner_begin, + int owner_count, + const torch::Tensor* communicator) const { + if (periodic) { + return deepmd::long_range::mesh_sum( + coord, charge, node_mask, n_node_per_frame, total_charge, box, + coefficient, kernel, channel, prefactor, order, spacing, owner_begin, + owner_count, communicator); + } + return deepmd::long_range::direct_sum( + coord, charge, node_mask, n_node_per_frame, total_charge, coefficient, + kernel, channel, prefactor, owner_begin, owner_count, communicator); + } +}; + +// Evaluates the frame of a system decomposed into the domains of +// ``communicator``, of which this rank holds one: its charged nodes are its +// owners, the charged nodes of every domain its partners. The nodes the rank +// does not own carry no charge, and the per-frame virial is the share of its +// owners. +ChargeSumOutputs evaluate_across_domains(const Evaluation& evaluation, + const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& communicator) { + TORCH_CHECK(n_node_per_frame.size(0) == 1, + "latent_charge_sum: an evaluation across MPI ranks takes one " + "frame, got ", + n_node_per_frame.size(0)); + const int64_t node_count = coord.size(0); + // The charged nodes of the frame; the padding nodes past it carry none. + // Their count is the one value the host has to wait for before the + // exchange. + const auto in_frame = + torch::arange(node_count, n_node_per_frame.options()) < n_node_per_frame; + const auto owned = torch::nonzero(node_mask & in_frame).squeeze(1); + const auto rows = + torch::cat({coord.index_select(0, owned), + charge.index_select(0, owned).to(torch::kFloat64)}, + 1) + .contiguous(); + const auto [partners, owner_begin] = + deepmd::long_range::gather_domain_rows(communicator, rows); + const int64_t partner_count = partners.size(0); + const auto partner_coord = partners.slice(1, 0, 3).contiguous(); + const auto partner_charge = + partners.slice(1, 3).to(torch::kFloat32).contiguous(); + const auto partner_mask = torch::ones({partner_count}, node_mask.options()); + const auto partner_frame = + torch::full({1}, partner_count, n_node_per_frame.options()); + const auto [owner_energy, owner_gradient, owner_force, owner_atom_virial, + virial] = + evaluation(partner_coord, partner_charge, partner_mask, partner_frame, + total_charge, static_cast(owner_begin), + static_cast(owned.size(0)), &communicator); + return {torch::zeros({node_count, 1}, coord.options()) + .index_copy_(0, owned, owner_energy), + torch::zeros({node_count, charge.size(1)}, charge.options()) + .index_copy_(0, owned, owner_gradient), + torch::zeros({node_count, 3}, charge.options()) + .index_copy_(0, owned, owner_force), + torch::zeros({node_count, 9}, charge.options()) + .index_copy_(0, owned, owner_atom_virial), + virial}; +} + +// Whether the frames carry cells: true when ``box`` is given and none of its +// rows is zero, false when it is absent or all of its rows are; a call that +// mixes the two is rejected. +bool periodic_frames(const c10::optional& box) { + if (!box.has_value()) { + return false; + } + const auto cells = box->contiguous(); + const double* value = cells.data_ptr(); + const int64_t frame_count = cells.size(0); + int64_t periodic = 0; + for (int64_t f = 0; f < frame_count; ++f) { + periodic += std::any_of(value + 9 * f, value + 9 * (f + 1), + [](double entry) { return entry != 0.0; }); + } + TORCH_CHECK(periodic == 0 || periodic == frame_count, + "latent_charge_sum: every frame of a call is periodic, or none " + "is"); + return periodic > 0; +} + +ChargeSumOutputs latent_charge_sum(torch::Tensor coord, + torch::Tensor charge, + torch::Tensor node_mask, + torch::Tensor n_node_per_frame, + torch::Tensor total_charge, + c10::optional box, + torch::Tensor coefficient, + torch::Tensor communicator, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing) { + TORCH_CHECK(coord.is_cuda() && charge.is_cuda() && node_mask.is_cuda() && + n_node_per_frame.is_cuda() && total_charge.is_cuda() && + coefficient.is_cuda(), + "latent_charge_sum: inputs must be CUDA tensors, except box and " + "communicator"); + TORCH_CHECK(charge.device() == coord.device() && + node_mask.device() == coord.device() && + n_node_per_frame.device() == coord.device() && + total_charge.device() == coord.device() && + coefficient.device() == coord.device(), + "latent_charge_sum: inputs must share one device"); + TORCH_CHECK(coord.dim() == 2 && coord.size(1) == 3 && + coord.scalar_type() == torch::kFloat64, + "latent_charge_sum: coord must be float64 with shape (N, 3)"); + TORCH_CHECK(charge.dim() == 2 && charge.size(0) == coord.size(0) && + charge.scalar_type() == torch::kFloat32, + "latent_charge_sum: charge must be float32 with shape (N, C)"); + TORCH_CHECK(node_mask.dim() == 1 && node_mask.size(0) == coord.size(0) && + node_mask.scalar_type() == torch::kBool, + "latent_charge_sum: node_mask must be bool with shape (N,)"); + TORCH_CHECK(n_node_per_frame.dim() == 1, + "latent_charge_sum: n_node_per_frame must be one-dimensional"); + TORCH_CHECK(total_charge.dim() == 1 && + total_charge.size(0) == n_node_per_frame.size(0) && + total_charge.scalar_type() == torch::kFloat64, + "latent_charge_sum: total_charge must be float64 with one entry " + "per frame"); + TORCH_CHECK( + !box.has_value() || + (box->device().is_cpu() && box->scalar_type() == torch::kFloat64 && + box->dim() == 2 && box->size(0) == n_node_per_frame.size(0) && + box->size(1) == 9), + "latent_charge_sum: box must be a float64 CPU tensor with shape " + "(F, 9)"); + TORCH_CHECK( + coefficient.dim() == 1 && coefficient.scalar_type() == torch::kFloat64, + "latent_charge_sum: coefficient must be a float64 vector"); + const int64_t channels = charge.size(1); + TORCH_CHECK(channels >= 1 && channels <= kMaxChannels, + "latent_charge_sum: the channel count must lie in [1, ", + kMaxChannels, "], got ", channels); + TORCH_CHECK(channel >= -1 && channel < channels, + "latent_charge_sum: the constrained channel must be -1 or lie " + "in [0, ", + channels, "), got ", channel); + if (kernel == kLes) { + TORCH_CHECK(coefficient.numel() == 1, + "latent_charge_sum: the erf kernel takes one coefficient"); + } else if (kernel == kSog) { + TORCH_CHECK(coefficient.numel() % 2 == 0 && coefficient.numel() > 0 && + coefficient.numel() <= 2 * kMaxGaussians, + "latent_charge_sum: the Gaussian kernel takes two " + "coefficients per Gaussian, for at most ", + kMaxGaussians, " Gaussians"); + } else { + TORCH_CHECK(false, "latent_charge_sum: unknown kernel ", kernel); + } + TORCH_CHECK(coord.size(0) < std::numeric_limits::max(), + "latent_charge_sum: too many nodes"); + + const c10::cuda::CUDAGuard guard(coord.device()); + coord = coord.contiguous(); + charge = charge.contiguous(); + node_mask = node_mask.contiguous(); + total_charge = total_charge.contiguous(); + const bool periodic = periodic_frames(box); + const torch::Tensor cells = periodic ? box->contiguous() : torch::Tensor(); + coefficient = coefficient.contiguous(); + const Evaluation evaluation{cells, coefficient, kernel, channel, + prefactor, order, spacing, periodic}; + if (deepmd::long_range::domain_of(communicator).size == 1) { + return evaluation(coord, charge, node_mask, n_node_per_frame, total_charge, + 0, static_cast(coord.size(0)), nullptr); + } + return evaluate_across_domains(evaluation, coord, charge, node_mask, + n_node_per_frame, total_charge, communicator); +} + +} // namespace + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.impl("latent_charge_sum", torch::kCUDA, &latent_charge_sum); +} diff --git a/source/op/pt/long_range/mesh_sum.cu b/source/op/pt/long_range/mesh_sum.cu new file mode 100644 index 0000000000..130cbdcb11 --- /dev/null +++ b/source/op/pt/long_range/mesh_sum.cu @@ -0,0 +1,1375 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA kernels of the latent-charge long range of periodic frames, by the +// smooth particle-mesh Ewald method (schema and formulas in ops.cc, the +// evaluation in deepmd/dpmodel/long_range/reciprocal.py). +// +// Every frame owns a mesh of n1 x n2 x n3 points spanning its cell. The +// charges spread to the mesh with cardinal B-splines of order p along the +// fractional coordinates s = r B^-1, the lattice vectors being the rows of the +// cell B: the scaled coordinate u_d = n_d s_d reaches the p points +// floor(u_d) - p + 1, ..., floor(u_d) with the weights M_p(u_d - m). A real +// FFT gives the half spectrum Q(m) of the charge mesh, and the influence +// function +// +// G(m) = prefactor * K^(k_m) / V * prod_d |D_d(m_d)|^-2, k_m = 2 pi B^-1 m, +// +// with the Fourier transform K^ of the kernel and the assignment factor +// D_d(m) = sum_l M_p(l + 1) exp(2 pi i m l / n_d), turns it into the +// potential mesh, the unnormalized inverse transform of G Q. The potential at +// a node is its interpolation by the B-splines of the node, less the potential +// prefactor * K(0) q of the charge at itself, and the force at fixed charges +// the negative gradient of that interpolation times the charges. The zero mode +// and the modes at the Nyquist frequency of an axis, which have no unique wave +// number, carry no weight. +// +// The virial of a frame is that of the strain of its cell, under which the +// fractional coordinates stay put while k_m and the volume change: +// +// W = sum_m e_m (1 + 2 dln K^ / d(k^2) k_m (x) k_m), e_m = G |Q|^2 / 2, +// +// the sum over the full spectrum, every interior plane of the half spectrum +// counting twice. Its per-node share is F_i (x) r_i, the force index first, +// plus the rest of W shared evenly over the charged nodes. +// +// The charges spread in 64-bit fixed point, so that the mesh, and every +// output, does not depend on the order of the atomic additions. The sums over +// the modes and over the nodes of a frame add float64 sums in a fixed order. +// The meshes, the transforms and the influence function are float32, against +// a float64 throughput of a small fraction of the float32 one on many +// devices; the B-spline weights are evaluated in float64, whose rounding to +// float32 leaves their sum unbiased. The forces then carry the rounding of +// float32, and the energy of a frame that of the reciprocal sum including the +// interaction of every charge with itself, which the self term cancels: a few +// 1e-7 of the self energy. +// +// The nodes whose outputs a call evaluates, its owners, are either every node +// or, across the domains of a decomposed system, the charged nodes of one MPI +// rank, whose frame then holds the charged nodes of every rank. Every rank +// spreads every charge and transforms the whole mesh, and interpolates only at +// its owners; the per-frame sums of the owners are added over the domains. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "charge_sum.h" +#include "frame_sums.cuh" + +namespace { + +constexpr int kThreads = 256; +constexpr int kMaxGaussians = 32; +constexpr int kMaxOrder = 12; + +enum KernelKind : int { kLes = 0, kSog = 1 }; + +// Resolution of the fixed-point charge mesh: 2^-40 e, far below the float32 +// precision of the raw charges, for a sum of at most 2^23 e at a point. +constexpr float kFixedScaleFloat = 1099511627776.f; // 2^40 + +constexpr double kPi = 3.14159265358979323846; + +// The cell of one frame and its mesh. +struct FrameMesh { + // B^-1, row-major: s_d = sum_j r_j inverse[3 j + d]. + double inverse[9]; + double volume; + // -n_d (B^-1)_jd, row-major in (j, d): the force at fixed charges from the + // derivatives of the energy with respect to the scaled coordinates u_d. + float gradient_to_force[9]; + int n[3]; + // The first point of the meshes of the frame, channel-major, in the mesh + // buffer of the call. + int64_t offset; +}; + +// The smallest even number of at least ``n`` without a prime factor above 7. +int64_t even_fft_size(int64_t n) { + int64_t m = std::max(2, n + n % 2); + while (true) { + int64_t rest = m; + for (const int64_t p : {2, 3, 5, 7}) { + while (rest % p == 0) { + rest /= p; + } + } + if (rest == 1) { + return m; + } + m += 2; + } +} + +FrameMesh frame_mesh(const double* cell, double spacing) { + FrameMesh mesh{}; + const double a = cell[0], b = cell[1], c = cell[2]; + const double d = cell[3], e = cell[4], f = cell[5]; + const double g = cell[6], h = cell[7], i = cell[8]; + const double det = + a * (e * i - f * h) - b * (d * i - f * g) + c * (d * h - e * g); + // A volume that is a vanishing fraction of the product of the lengths of + // the lattice vectors, or not finite, spans no lattice. + const double lengths = std::sqrt(a * a + b * b + c * c) * + std::sqrt(d * d + e * e + f * f) * + std::sqrt(g * g + h * h + i * i); + TORCH_CHECK(std::abs(det) > 1e-6 * lengths, + "latent_charge_sum: a periodic cell is singular"); + const double adjugate[9] = {e * i - f * h, c * h - b * i, b * f - c * e, + f * g - d * i, a * i - c * g, c * d - a * f, + d * h - e * g, b * g - a * h, a * e - b * d}; + for (int k = 0; k < 9; ++k) { + mesh.inverse[k] = adjugate[k] / det; + } + mesh.volume = std::abs(det); + // The squared lengths of the lattice vectors are formed term by term, as + // the Python front end forms them, so that both round the mesh size + // identically. + for (int axis = 0; axis < 3; ++axis) { + const double* v = cell + 3 * axis; + const double length = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + const double points = std::ceil(length / spacing); + TORCH_CHECK(points < 1 << 20, + "latent_charge_sum: the mesh of a cell exceeds 2^20 points " + "along a lattice vector"); + mesh.n[axis] = + static_cast(even_fft_size(static_cast(points))); + } + TORCH_CHECK(int64_t{mesh.n[0]} * mesh.n[1] * mesh.n[2] < + std::numeric_limits::max(), + "latent_charge_sum: the mesh of a cell exceeds 2^31 points"); + for (int j = 0; j < 3; ++j) { + for (int d = 0; d < 3; ++d) { + mesh.gradient_to_force[3 * j + d] = + static_cast(-mesh.n[d] * mesh.inverse[3 * j + d]); + } + } + return mesh; +} + +__host__ __device__ int64_t mesh_points(const FrameMesh& mesh) { + return int64_t{mesh.n[0]} * mesh.n[1] * mesh.n[2]; +} + +// Nodes per block of the kernels that spread and interpolate, one warp each. +constexpr int kStencilWarps = 4; +constexpr int kWarp = 32; + +// The stencil of one node along one axis of its mesh: for each of its p +// points on the axis, its offset on the flat mesh of a frame along that axis, +// the index of the point times the stride of the axis, and its B-spline +// weight and slope, evaluated in float64 and stored in float32, the precision +// of the outputs: rounding a recursion of float32 steps instead would bias +// the sum of the weights. +template +struct AxisStencil { + int offset[kOrder]; + float value[kOrder]; + float slope[kOrder]; +}; + +// The stencil of a node along axis ``d``. The scaled coordinate u = n_d s_d +// of the fractional coordinate s_d reaches the points floor(u) - p + 1 + k, +// k = 0, ..., p - 1, with the weights value[k] = M_p(t + p - 1 - k), t the +// fractional part of u, and slope[k] their derivative with respect to u. The +// recursion M_n(t + j) = ((t + j) M_{n-1}(t + j) + (n - t - j) M_{n-1}(t + j - +// 1)) / (n - 1) builds a_j = M_n(t + j) from M_2; the slope follows from +// M_p'(x) = M_{p-1}(x) - M_{p-1}(x - 1). +template +__device__ void axis_stencil(const double* coord, + const FrameMesh& mesh, + int d, + AxisStencil& stencil) { + double s = coord[0] * mesh.inverse[d] + coord[1] * mesh.inverse[3 + d] + + coord[2] * mesh.inverse[6 + d]; + s -= floor(s); + const double u = s * mesh.n[d]; + const double base = floor(u); + const double t = u - base; + double a[kOrder]; + a[0] = t; + a[1] = 1.0 - t; +#pragma unroll + for (int n = 3; n <= kOrder; ++n) { + if (n == kOrder) { + // The weights of order p - 1 give the slopes. +#pragma unroll + for (int j = 0; j < kOrder; ++j) { + const double high = j < kOrder - 1 ? a[j] : 0.0; + const double low = j > 0 ? a[j - 1] : 0.0; + stencil.slope[kOrder - 1 - j] = static_cast(high - low); + } + } + const double scale = 1.0 / (n - 1); + a[n - 1] = scale * (1.0 - t) * a[n - 2]; +#pragma unroll + for (int j = n - 2; j > 0; --j) { + a[j] = scale * ((t + j) * a[j] + (n - t - j) * a[j - 1]); + } + a[0] = scale * t * a[0]; + } + if constexpr (kOrder == 2) { + stencil.slope[0] = -1.f; + stencil.slope[1] = 1.f; + } + const int n = mesh.n[d]; + const int stride = d == 0 ? mesh.n[1] * mesh.n[2] : d == 1 ? mesh.n[2] : 1; + // A stencil wider than the axis wraps more than once. + const int first = ((static_cast(base) - kOrder + 1) % n + n) % n; +#pragma unroll + for (int k = 0; k < kOrder; ++k) { + stencil.value[k] = static_cast(a[kOrder - 1 - k]); + stencil.offset[k] = (first + k) % n * stride; + } +} + +// Writes, one thread per node and axis, the stencils of the charged nodes +// along the three axes of their meshes, and, on the thread of the first axis, +// the charges of every node, shifted in the constrained channel and zero on +// the nodes that carry none. Every float64 warp instruction occupies the +// float64 pipe of the device whatever the number of its active lanes, so the +// float64 work of a node runs here with every lane busy rather than in the +// warps of the kernels that spread and interpolate. +template +__global__ void __launch_bounds__(kThreads) + node_stencils(const double* __restrict__ coord, + const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + const FrameMesh* __restrict__ frames, + const double* __restrict__ frame_stat, + int frame_count, + int64_t node_count, + int channels, + int channel, + AxisStencil* __restrict__ stencils, + float* __restrict__ node_charge) { + for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; + index < 3 * node_count; index += int64_t{gridDim.x} * blockDim.x) { + const int64_t node = index / 3; + const int axis = static_cast(index - 3 * node); + const int f = frame_of(frame_ptr, frame_count, node); + const bool charged = f >= 0 && node_mask[node]; + if (charged) { + axis_stencil(coord + 3 * node, frames[f], axis, stencils[index]); + } + if (axis == 0) { + const double shift = charged && channel >= 0 + ? frame_stat[f * kStatWidth + kStatShift] + : 0.0; + for (int c = 0; c < channels; ++c) { + const double raw = charged ? charge[node * channels + c] : 0.0; + node_charge[node * channels + c] = + static_cast(c == channel ? raw - shift : raw); + } + } + } +} + +// Rounds ``x`` to the nearest integer, ties to even, as __float2ll_rn does, +// in integer operations: the conversion instruction occupies the float64 pipe. +// A finite x is a 24-bit significand times 2^exponent; a negative exponent +// shifts the significand right with rounding. Magnitudes of 2^63 and beyond, +// infinities and NaN saturate. +__device__ __forceinline__ long long float_to_fixed(float x) { + const unsigned bits = __float_as_uint(x); + const int biased = static_cast((bits >> 23) & 0xffu); + // |x| < 1/2, zero and subnormal numbers included, rounds to zero. + if (biased < 126) { + return 0; + } + const int exponent = biased - 150; + if (exponent >= 40) { + return bits >> 31 ? std::numeric_limits::min() + : std::numeric_limits::max(); + } + const long long significand = (bits & 0x7fffffu) | 0x800000u; + long long magnitude; + if (exponent >= 0) { + magnitude = significand << exponent; + } else { + const int shift = -exponent; + const long long kept = significand >> shift; + const long long rest = significand & ((1ll << shift) - 1); + const long long half = 1ll << (shift - 1); + magnitude = kept + (rest > half || (rest == half && (kept & 1)) ? 1 : 0); + } + return bits >> 31 ? -magnitude : magnitude; +} + +// The assignment of the p^3 points of a stencil to the lanes of a warp: the +// lanes take the (k2, k3) pairs, k3 running fastest, so that a warp touches +// contiguous runs of the mesh, and every lane walks the k1 axis for its +// pairs; while the pairs are fewer than the lanes, groups of lanes split the +// k1 axis. +template +struct LaneSlots { + static constexpr int kPairs = kOrder * kOrder; + static constexpr int kGroups = kPairs >= kWarp ? 1 : kWarp / kPairs; + static constexpr int kSlots = kGroups * kPairs; + static constexpr int kRounds = (kSlots + kWarp - 1) / kWarp; +}; + +// Copies the stencils of a node along its three axes to the shared memory of +// its warp. +template +__device__ __forceinline__ void load_stencil( + const AxisStencil* __restrict__ stencils, + int64_t node, + int lane, + AxisStencil* shared) { + constexpr int kWords = 3 * sizeof(AxisStencil) / sizeof(int); + const int* source = reinterpret_cast(stencils + 3 * node); + int* target = reinterpret_cast(shared); + for (int word = lane; word < kWords; word += kWarp) { + target[word] = source[word]; + } + __syncwarp(); +} + +// Spreads the charges of the charged nodes to the meshes of their frames, one +// warp per node, one pass per channel. +template +__global__ void __launch_bounds__(kStencilWarps* kWarp) + spread_charges(const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + const FrameMesh* __restrict__ frames, + const AxisStencil* __restrict__ stencils, + const float* __restrict__ node_charge, + int frame_count, + int64_t node_count, + int channels, + unsigned long long* __restrict__ mesh) { + using Slots = LaneSlots; + __shared__ AxisStencil shared[kStencilWarps][3]; + const int warp = threadIdx.x / kWarp; + const int lane = threadIdx.x % kWarp; + const AxisStencil* axes = shared[warp]; + for (int64_t node = int64_t{blockIdx.x} * kStencilWarps + warp; + node < node_count; node += int64_t{gridDim.x} * kStencilWarps) { + const int f = frame_of(frame_ptr, frame_count, node); + if (f < 0 || !node_mask[node]) { + continue; + } + load_stencil(stencils, node, lane, shared[warp]); + const FrameMesh& grid = frames[f]; + for (int c = 0; c < channels; ++c) { + const float q = node_charge[node * channels + c] * kFixedScaleFloat; + unsigned long long* channel_mesh = + mesh + grid.offset + c * mesh_points(grid); +#pragma unroll + for (int round = 0; round < Slots::kRounds; ++round) { + const int slot = lane + kWarp * round; + if (slot < Slots::kSlots) { + const int group = slot / Slots::kPairs; + const int pair = slot - group * Slots::kPairs; + const int k2 = pair / kOrder; + const int k3 = pair - k2 * kOrder; + const float weight = q * axes[1].value[k2] * axes[2].value[k3]; + const int column = axes[1].offset[k2] + axes[2].offset[k3]; +#pragma unroll + for (int k1 = group; k1 < kOrder; k1 += Slots::kGroups) { + atomicAdd(channel_mesh + axes[0].offset[k1] + column, + static_cast( + float_to_fixed(weight * axes[0].value[k1]))); + } + } + } + } + // The stencil of the next node of the warp overwrites this one. + __syncwarp(); + } +} + +// Converts the fixed-point mesh to the float32 mesh that the transforms read. +__global__ void fixed_to_float(int64_t count, + const unsigned long long* __restrict__ fixed, + float* __restrict__ mesh) { + constexpr float kFixedInverse = 1.f / kFixedScaleFloat; + for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; + index < count; index += int64_t{gridDim.x} * blockDim.x) { + mesh[index] = + __ll2float_rn(static_cast(fixed[index])) * kFixedInverse; + } +} + +// The weights M_p(1), ..., M_p(p - 1) of the assignment factor. +struct IntegerWeights { + double value[kMaxOrder - 1]; +}; + +// The inverse squared assignment factor |D_d(m)|^-2 along every axis of the +// half spectrum of one frame: n1 entries, then n2, then n3 / 2 + 1. +__global__ void assignment_factor(FrameMesh mesh, + int order, + IntegerWeights integer_weight, + float* __restrict__ factor) { + const int half = mesh.n[2] / 2 + 1; + const int total = mesh.n[0] + mesh.n[1] + half; + for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < total; + index += gridDim.x * blockDim.x) { + int axis = 0; + int m = index; + if (m >= mesh.n[0]) { + m -= mesh.n[0]; + axis = 1; + if (m >= mesh.n[1]) { + m -= mesh.n[1]; + axis = 2; + } + } + const double theta = 2.0 * kPi * m / mesh.n[axis]; + double re = 0.0; + double im = 0.0; + for (int l = 0; l < order - 1; ++l) { + double s; + double c; + sincos(theta * l, &s, &c); + re += integer_weight.value[l] * c; + im += integer_weight.value[l] * s; + } + factor[index] = static_cast(1.0 / (re * re + im * im)); + } +} + +// The terms of the transform of the kernel, in float32, the precision of the +// transforms: 1 / (4 alpha^2) of the les kernel, or, for every Gaussian of the +// sog kernel, its weight a_g (2 pi)^(3/2) sigma_g^3 and half its squared +// width, from ``coefficient``, which holds alpha, or the amplitudes and then +// the rates c_g = log2(e) / (2 sigma_g^2). +struct KernelTerms { + float inv_four_alpha_sq; + float weight[kMaxGaussians]; + float half_width_sq[kMaxGaussians]; +}; + +__device__ void load_kernel_terms(int kernel, + const double* __restrict__ coefficient, + int gaussians, + KernelTerms& terms) { + if (kernel == kLes) { + if (threadIdx.x == 0) { + const double alpha = coefficient[0]; + terms.inv_four_alpha_sq = static_cast(0.25 / (alpha * alpha)); + } + return; + } + // (2 pi)^{3/2} + constexpr double kGaussianNorm = 15.749609945722419; + for (int g = threadIdx.x; g < gaussians; g += blockDim.x) { + // sigma^2 = log2(e) / (2 c) = 1 / (2 c ln 2) + const double width_sq = 0.72134752044448170 / coefficient[gaussians + g]; + terms.weight[g] = static_cast(coefficient[g] * kGaussianNorm * + width_sq * sqrt(width_sq)); + terms.half_width_sq[g] = static_cast(0.5 * width_sq); + } +} + +// The transform of the kernel at k^2 and its logarithmic derivative with +// respect to k^2. +__device__ __forceinline__ void kernel_transform(int kernel, + const KernelTerms& terms, + int gaussians, + float k_sq, + float& value, + float& log_slope) { + if (kernel == kLes) { + value = 4.f * static_cast(kPi) * + expf(-k_sq * terms.inv_four_alpha_sq) / k_sq; + log_slope = -terms.inv_four_alpha_sq - 1.f / k_sq; + return; + } + value = 0.f; + float slope = 0.f; + for (int g = 0; g < gaussians; ++g) { + const float term = terms.weight[g] * expf(-terms.half_width_sq[g] * k_sq); + value += term; + slope -= terms.half_width_sq[g] * term; + } + log_slope = value != 0.f ? slope / value : 0.f; +} + +// The influence function of one frame on its half spectrum: multiplies every +// mode of the spectra of the channels by G in place and sums, per block, the +// energy of the modes and their virial, xx yy zz xy xz yz. The modes are +// evaluated in float32, the precision of the transforms, summed in float32 +// over the modes of a row that a lane takes, and in float64 over the rows. +constexpr int kModeSums = 7; + +__global__ void __launch_bounds__(kThreads) + apply_influence(FrameMesh mesh, + int channels, + int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor, + const float* __restrict__ factor, + float2* __restrict__ spectrum, + double* __restrict__ block_sum) { + __shared__ KernelTerms terms; + load_kernel_terms(kernel, coefficient, gaussians, terms); + __syncthreads(); + const int n1 = mesh.n[0]; + const int n2 = mesh.n[1]; + const int n3 = mesh.n[2]; + const int half = n3 / 2 + 1; + const int modes = n1 * n2 * half; + // k_a = 2 pi sum_d (B^-1)_ad m_d + float reciprocal[9]; +#pragma unroll + for (int k = 0; k < 9; ++k) { + reciprocal[k] = static_cast(2.0 * kPi * mesh.inverse[k]); + } + const float scale = static_cast(prefactor / mesh.volume); + // Every warp takes rows (i1, i2) of the half spectrum, its lanes the modes + // of a row, which lie contiguously. + const int lane = threadIdx.x % kWarp; + const int warps = gridDim.x * (kThreads / kWarp); + double sum[kModeSums] = {}; + for (int row = blockIdx.x * (kThreads / kWarp) + threadIdx.x / kWarp; + row < n1 * n2; row += warps) { + float row_sum[kModeSums] = {}; + const int i1 = row / n2; + const int i2 = row - i1 * n2; + const int m1 = i1 <= n1 / 2 ? i1 : i1 - n1; + const int m2 = i2 <= n2 / 2 ? i2 : i2 - n2; + const bool row_weightless = 2 * m1 == n1 || 2 * m2 == n2; + float k_row[3]; +#pragma unroll + for (int a = 0; a < 3; ++a) { + k_row[a] = reciprocal[3 * a] * m1 + reciprocal[3 * a + 1] * m2; + } + const float row_factor = scale * factor[i1] * factor[n1 + i2]; + for (int i3 = lane; i3 < half; i3 += kWarp) { + const int m3 = i3; + const bool weightless = + row_weightless || 2 * m3 == n3 || (m1 == 0 && m2 == 0 && m3 == 0); + float influence = 0.f; + float k[3] = {0.f, 0.f, 0.f}; + float log_slope = 0.f; + if (!weightless) { +#pragma unroll + for (int a = 0; a < 3; ++a) { + k[a] = k_row[a] + reciprocal[3 * a + 2] * m3; + } + const float k_sq = k[0] * k[0] + k[1] * k[1] + k[2] * k[2]; + float transform; + kernel_transform(kernel, terms, gaussians, k_sq, transform, log_slope); + influence = row_factor * transform * factor[n1 + n2 + i3]; + } + // An interior plane of the half spectrum stands for its mirror as well. + const float multiplicity = m3 == 0 ? 1.f : 2.f; + const int index = row * half + i3; + float energy = 0.f; + for (int c = 0; c < channels; ++c) { + float2& value = spectrum[int64_t{c} * modes + index]; + energy += value.x * value.x + value.y * value.y; + value.x *= influence; + value.y *= influence; + } + energy *= 0.5f * multiplicity * influence; + const float twice = 2.f * log_slope * energy; + row_sum[0] += energy; + row_sum[1] += energy + twice * k[0] * k[0]; + row_sum[2] += energy + twice * k[1] * k[1]; + row_sum[3] += energy + twice * k[2] * k[2]; + row_sum[4] += twice * k[0] * k[1]; + row_sum[5] += twice * k[0] * k[2]; + row_sum[6] += twice * k[1] * k[2]; + } +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + sum[s] += row_sum[s]; + } + } + __shared__ double partial[kModeSums][kThreads]; + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] = sum[s]; + } + __syncthreads(); + for (int stride = kThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] += partial[s][threadIdx.x + stride]; + } + } + __syncthreads(); + } + if (threadIdx.x < kModeSums) { + block_sum[blockIdx.x * kModeSums + threadIdx.x] = partial[threadIdx.x][0]; + } +} + +// Adds the block sums of apply_influence into the virial of the frame, +// row-major: every thread sums a fixed strided subset of the blocks, and the +// block adds the thread sums in a fixed order. +__global__ void __launch_bounds__(kThreads) + mode_virial(const double* __restrict__ block_sum, + int block_count, + double* __restrict__ virial) { + double value[kModeSums] = {}; + for (int block = threadIdx.x; block < block_count; block += kThreads) { +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + value[s] += block_sum[block * kModeSums + s]; + } + } + __shared__ double partial[kModeSums][kThreads]; +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] = value[s]; + } + __syncthreads(); + for (int stride = kThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] += partial[s][threadIdx.x + stride]; + } + } + __syncthreads(); + } + // Components 1..6 are xx yy zz xy xz yz; the energy is not stored. + constexpr int kEntry[kModeSums][2] = {{-1, -1}, {0, 0}, {4, 4}, {8, 8}, + {1, 3}, {2, 6}, {5, 7}}; + if (threadIdx.x > 0 && threadIdx.x < kModeSums) { + virial[kEntry[threadIdx.x][0]] = partial[threadIdx.x][0]; + virial[kEntry[threadIdx.x][1]] = partial[threadIdx.x][0]; + } +} + +// The contact value K(0) of the kernel. +__device__ double contact_value(int kernel, + const double* __restrict__ coefficient, + int gaussians) { + if (kernel == kLes) { + return 1.1283791670955126 * coefficient[0]; + } + double value = 0.0; + for (int g = 0; g < gaussians; ++g) { + value += coefficient[g]; + } + return value; +} + +// Adds ``value`` over the lanes of a warp in a fixed butterfly order; every +// lane receives the total. +__device__ __forceinline__ float warp_sum(float value) { +#pragma unroll + for (int stride = kWarp / 2; stride > 0; stride /= 2) { + value += __shfl_xor_sync(0xffffffffu, value, stride); + } + return value; +} + +// Interpolates the potential meshes at every owner node, one warp per node, +// one pass per channel: the potential phi^rec_c of the mesh, which still holds +// that of the charge at itself, and the force at fixed charges. Owners that +// carry no charge get zeros. Every lane sums its points in float32, the +// precision of the meshes, and the warp adds the lane sums. +template +__global__ void __launch_bounds__(kStencilWarps* kWarp) + interpolate_owners(const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + const FrameMesh* __restrict__ frames, + const AxisStencil* __restrict__ stencils, + const float* __restrict__ node_charge, + const float* __restrict__ potential_mesh, + int frame_count, + int owner_begin, + int owner_count, + int channels, + float* __restrict__ potential, + float* __restrict__ force) { + using Slots = LaneSlots; + __shared__ AxisStencil shared[kStencilWarps][3]; + const int warp = threadIdx.x / kWarp; + const int lane = threadIdx.x % kWarp; + const AxisStencil* axes = shared[warp]; + for (int owner = blockIdx.x * kStencilWarps + warp; owner < owner_count; + owner += gridDim.x * kStencilWarps) { + const int64_t node = int64_t{owner_begin} + owner; + const int f = frame_of(frame_ptr, frame_count, node); + if (f < 0 || !node_mask[node]) { + for (int c = lane; c < channels; c += kWarp) { + potential[int64_t{owner} * channels + c] = 0.f; + } + if (lane < 3) { + force[int64_t{owner} * 3 + lane] = 0.f; + } + continue; + } + load_stencil(stencils, node, lane, shared[warp]); + const FrameMesh& grid = frames[f]; + // The derivative of sum_c q_c phi_c with respect to u_d. + float gradient[3] = {0.f, 0.f, 0.f}; + for (int c = 0; c < channels; ++c) { + const float* channel_mesh = + potential_mesh + grid.offset + c * mesh_points(grid); + float phi = 0.f; + float slope[3] = {0.f, 0.f, 0.f}; +#pragma unroll + for (int round = 0; round < Slots::kRounds; ++round) { + const int slot = lane + kWarp * round; + if (slot < Slots::kSlots) { + const int group = slot / Slots::kPairs; + const int pair = slot - group * Slots::kPairs; + const int k2 = pair / kOrder; + const int k3 = pair - k2 * kOrder; + const float v2 = axes[1].value[k2]; + const float v3 = axes[2].value[k3]; + const float v23 = v2 * v3; + const float s2v3 = axes[1].slope[k2] * v3; + const float v2s3 = v2 * axes[2].slope[k3]; + const int column = axes[1].offset[k2] + axes[2].offset[k3]; +#pragma unroll + for (int k1 = group; k1 < kOrder; k1 += Slots::kGroups) { + const float value = channel_mesh[axes[0].offset[k1] + column]; + const float v1_value = axes[0].value[k1] * value; + phi += v1_value * v23; + slope[0] += axes[0].slope[k1] * value * v23; + slope[1] += v1_value * s2v3; + slope[2] += v1_value * v2s3; + } + } + } + phi = warp_sum(phi); + const float q = node_charge[node * channels + c]; +#pragma unroll + for (int d = 0; d < 3; ++d) { + gradient[d] += q * warp_sum(slope[d]); + } + if (lane == 0) { + potential[int64_t{owner} * channels + c] = phi; + } + } + if (lane < 3) { + // F_j = -sum_d dE/du_d n_d (B^-1)_jd + force[int64_t{owner} * 3 + lane] = + gradient[0] * grid.gradient_to_force[3 * lane] + + gradient[1] * grid.gradient_to_force[3 * lane + 1] + + gradient[2] * grid.gradient_to_force[3 * lane + 2]; + } + // The stencil of the next node of the warp overwrites this one. + __syncwarp(); + } +} + +// The potential prefactor * K(0) of a unit charge at itself, which the +// potential mesh holds and the potential of a node excludes, for every thread +// of the block. +__device__ double self_potential(int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor) { + __shared__ double self; + if (threadIdx.x == 0) { + self = prefactor * contact_value(kernel, coefficient, gaussians); + } + __syncthreads(); + return self; +} + +// Per-frame sums over one slice of the owners: the potential of the charged +// owners in the constrained channel, the number of charged owners and the +// nine components of sum_i F_i (x) r_i. +constexpr int kOwnerSums = 11; + +__global__ void __launch_bounds__(kFrameThreads) + owner_slices(const float* __restrict__ potential, + const float* __restrict__ force, + const double* __restrict__ owner_coord, + const float* __restrict__ owner_charge, + const bool* __restrict__ owner_mask, + const int64_t* __restrict__ owner_frame_ptr, + int channels, + int channel, + int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor, + double* __restrict__ slice_sum) { + const double self = self_potential(kernel, coefficient, gaussians, prefactor); + const NodeRange slice = slice_of_frame(owner_frame_ptr); + double value[kOwnerSums] = {}; + for (int64_t owner = slice.begin + threadIdx.x; owner < slice.end; + owner += kFrameThreads) { + if (!owner_mask[owner]) { + continue; + } + if (channel >= 0) { + value[0] += double{potential[owner * channels + channel]} - + self * owner_charge[owner * channels + channel]; + } + value[1] += 1.0; + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + value[2 + 3 * a + b] += + double{force[owner * 3 + a]} * owner_coord[owner * 3 + b]; + } + } + } + write_slice_sum(value, slice_sum); +} + +// Adds the slice sums of owner_slices of every frame in slice order. +__global__ void owner_totals(const double* __restrict__ slice_sum, + int slice_count, + double* __restrict__ total) { + if (threadIdx.x >= kOwnerSums) { + return; + } + double value = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + value += + slice_sum[(static_cast(blockIdx.x) * slice_count + slice) * + kOwnerSums + + threadIdx.x]; + } + total[blockIdx.x * kOwnerSums + threadIdx.x] = value; +} + +// Writes the outputs of every owner: the atomic energy q . phi / 2 with phi = +// phi^rec - prefactor K(0) q, the charge gradient w (phi - [c = ch] mean), and +// the virial F (x) r plus the share of the rest of the virial W of the frame +// over its charged nodes. The first threads write the share of the virial of +// every frame held by the owners, sum_i F_i (x) r_i + n (W - sum_j F_j (x) +// r_j) / N over the n charged owners of the N charged nodes of the frame. +// ``owner_local`` holds the owner sums of this domain, ``owner_total`` those +// of every domain. +__global__ void finish_owners(const float* __restrict__ potential, + const float* __restrict__ force, + const double* __restrict__ owner_coord, + const float* __restrict__ owner_charge, + const bool* __restrict__ owner_mask, + const int64_t* __restrict__ owner_frame_ptr, + const double* __restrict__ frame_virial, + const double* __restrict__ owner_local, + const double* __restrict__ owner_total, + const double* __restrict__ frame_stat, + int frame_count, + int owner_count, + int channels, + int channel, + int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor, + double* __restrict__ atom_energy, + float* __restrict__ charge_gradient, + float* __restrict__ atom_virial, + float* __restrict__ virial) { + const double self = self_potential(kernel, coefficient, gaussians, prefactor); + const int start = blockIdx.x * blockDim.x + threadIdx.x; + const int threads = gridDim.x * blockDim.x; + for (int index = start; index < 9 * frame_count; index += threads) { + const int f = index / 9; + const int k = index - 9 * f; + const double count = frame_stat[f * kStatWidth + kStatCount]; + const double* local = owner_local + f * kOwnerSums; + const double rest = + count > 0.0 + ? (frame_virial[index] - owner_total[f * kOwnerSums + 2 + k]) / + count + : 0.0; + virial[index] = static_cast(local[2 + k] + local[1] * rest); + } + for (int owner = start; owner < owner_count; owner += threads) { + const int f = frame_of(owner_frame_ptr, frame_count, owner); + const bool charged = f >= 0 && owner_mask[owner]; + // The charged nodes of the frame, of every domain. + const double count = + charged ? frame_stat[f * kStatWidth + kStatCount] : 1.0; + const double mean = + charged && channel >= 0 ? owner_total[f * kOwnerSums] / count : 0.0; + double energy = 0.0; + for (int c = 0; c < channels; ++c) { + const int64_t entry = int64_t{owner} * channels + c; + const double q = owner_charge[entry]; + const double phi = charged ? potential[entry] - self * q : 0.0; + energy += q * phi; + charge_gradient[entry] = + static_cast(c == channel ? phi - mean : phi); + } + atom_energy[owner] = 0.5 * energy; + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + const int k = 3 * a + b; + const double moment = double{force[int64_t{owner} * 3 + a]} * + owner_coord[int64_t{owner} * 3 + b]; + const double rest = charged ? (frame_virial[f * 9 + k] - + owner_total[f * kOwnerSums + 2 + k]) / + count + : 0.0; + atom_virial[int64_t{owner} * 9 + k] = static_cast(moment + rest); + } + } + } +} + +IntegerWeights integer_bspline_weights(int order) { + IntegerWeights weight{}; + for (int l = 0; l < order - 1; ++l) { + const double x = l + 1.0; + double value = 0.0; + double binomial = 1.0; + for (int j = 0; j <= order; ++j) { + if (j > 0) { + binomial = binomial * (order - j + 1) / j; + } + const double base = std::max(x - j, 0.0); + value += (j % 2 == 0 ? 1.0 : -1.0) * binomial * std::pow(base, order - 1); + } + double factorial = 1.0; + for (int k = 2; k < order; ++k) { + factorial *= k; + } + weight.value[l] = value / factorial; + } + return weight; +} + +void check_cufft(cufftResult result, const char* operation) { + TORCH_CHECK(result == CUFFT_SUCCESS, "latent_charge_sum: cuFFT ", operation, + " failed with error ", static_cast(result)); +} + +// The real-to-complex and complex-to-real transforms of the meshes of one +// shape, batched over the channels. The plans allocate no work area: every +// transform runs on a work area of the workspace of its call, set with the +// stream under ``mutex`` just before the transform is enqueued, so that +// concurrent calls never share one. +struct MeshPlans { + cufftHandle forward = 0; + cufftHandle inverse = 0; + size_t work_bytes = 0; + c10::DeviceIndex device = 0; + std::mutex mutex; + + MeshPlans() = default; + MeshPlans(const MeshPlans&) = delete; + MeshPlans& operator=(const MeshPlans&) = delete; + // A plan leaves the cache only past kPlanCacheSize shapes; its transforms + // may still be in flight, so the device finishes them first. The + // destructor reports no error. + ~MeshPlans() { + int current = 0; + cudaGetDevice(¤t); + cudaSetDevice(device); + cudaDeviceSynchronize(); + if (forward != 0) { + cufftDestroy(forward); + } + if (inverse != 0) { + cufftDestroy(inverse); + } + cudaSetDevice(current); + } +}; + +// The plans of the meshes of ``mesh`` with ``channels`` channels on the +// current device, kept across calls; beyond kPlanCacheSize shapes, the least +// recently used are released. The cache itself lives until the process ends +// and is never destroyed, since its plans would otherwise be released after +// the CUDA runtime shuts down. +constexpr size_t kPlanCacheSize = 16; + +std::shared_ptr mesh_plans(const FrameMesh& mesh, int channels) { + using Key = std::array; + using Cache = std::list>>; + static std::mutex cache_mutex; + static Cache& cache = *new Cache(); + const Key key = {c10::cuda::current_device(), mesh.n[0], mesh.n[1], mesh.n[2], + channels}; + // A plan that leaves the cache is released after the lock, since its + // release waits for the device. + std::shared_ptr evicted; + const std::lock_guard lock(cache_mutex); + for (auto entry = cache.begin(); entry != cache.end(); ++entry) { + if (entry->first == key) { + cache.splice(cache.begin(), cache, entry); + return entry->second; + } + } + auto plans = std::make_shared(); + plans->device = c10::cuda::current_device(); + long long shape[3] = {mesh.n[0], mesh.n[1], mesh.n[2]}; + const long long points = shape[0] * shape[1] * shape[2]; + const long long modes = shape[0] * shape[1] * (shape[2] / 2 + 1); + size_t forward_bytes = 0; + size_t inverse_bytes = 0; + check_cufft(cufftCreate(&plans->forward), "plan creation"); + check_cufft(cufftSetAutoAllocation(plans->forward, 0), "plan setup"); + check_cufft( + cufftMakePlanMany64(plans->forward, 3, shape, nullptr, 1, points, nullptr, + 1, modes, CUFFT_R2C, channels, &forward_bytes), + "plan creation"); + check_cufft(cufftCreate(&plans->inverse), "plan creation"); + check_cufft(cufftSetAutoAllocation(plans->inverse, 0), "plan setup"); + check_cufft( + cufftMakePlanMany64(plans->inverse, 3, shape, nullptr, 1, modes, nullptr, + 1, points, CUFFT_C2R, channels, &inverse_bytes), + "plan creation"); + plans->work_bytes = std::max(forward_bytes, inverse_bytes); + cache.emplace_front(key, plans); + if (cache.size() > kPlanCacheSize) { + evicted = std::move(cache.back().second); + cache.pop_back(); + } + return plans; +} + +int blocks_for(int64_t work) { + return static_cast(std::max( + 1, std::min((work + kThreads - 1) / kThreads, 65535))); +} + +// Blocks of a stencil kernel over ``nodes`` nodes, one warp each. +int stencil_blocks_for(int64_t nodes) { + return static_cast(std::max( + 1, std::min((nodes + kStencilWarps - 1) / kStencilWarps, + int64_t{1} << 20))); +} + +// Calls ``launch`` with the B-spline order as a compile-time constant, for +// which the stencil kernels unroll their loops. +template +void dispatch_order(int64_t order, const Launch& launch) { + switch (order) { + case 2: + return launch(std::integral_constant{}); + case 3: + return launch(std::integral_constant{}); + case 4: + return launch(std::integral_constant{}); + case 5: + return launch(std::integral_constant{}); + case 6: + return launch(std::integral_constant{}); + case 7: + return launch(std::integral_constant{}); + case 8: + return launch(std::integral_constant{}); + case 9: + return launch(std::integral_constant{}); + case 10: + return launch(std::integral_constant{}); + case 11: + return launch(std::integral_constant{}); + case 12: + return launch(std::integral_constant{}); + default: + TORCH_CHECK(false, "latent_charge_sum: unsupported B-spline order ", + order); + } +} + +// The value of ``evaluate`` at the B-spline order as a compile-time constant. +template +auto dispatch_order_value(int64_t order, const Evaluate& evaluate) { + decltype(evaluate(std::integral_constant{})) value{}; + dispatch_order( + order, [&](auto order_constant) { value = evaluate(order_constant); }); + return value; +} + +} // namespace + +namespace deepmd { +namespace long_range { + +ChargeSumOutputs mesh_sum(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& box, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing, + int owner_begin, + int owner_count, + const torch::Tensor* communicator) { + TORCH_CHECK(order >= 2 && order <= kMaxOrder, + "latent_charge_sum: the B-spline order must lie in [2, ", + kMaxOrder, "], got ", order); + TORCH_CHECK(std::isfinite(spacing) && spacing > 0.0, + "latent_charge_sum: the mesh spacing must be positive, got ", + spacing); + const int64_t node_count = coord.size(0); + TORCH_CHECK(node_count < std::numeric_limits::max(), + "latent_charge_sum: too many nodes"); + const int frame_count = static_cast(n_node_per_frame.size(0)); + const int channels = static_cast(charge.size(1)); + const int gaussians = + kernel == kSog ? static_cast(coefficient.numel() / 2) : 0; + const auto double_options = coord.options(); + const auto float_options = charge.options(); + if (node_count == 0 || frame_count == 0) { + return {torch::zeros({owner_count, 1}, double_options), + torch::zeros({owner_count, channels}, float_options), + torch::zeros({owner_count, 3}, float_options), + torch::zeros({owner_count, 9}, float_options), + torch::zeros({frame_count, 9}, float_options)}; + } + const bool every_node = owner_begin == 0 && owner_count == node_count; + TORCH_CHECK(every_node || frame_count == 1, + "latent_charge_sum: owners other than every node lie in one " + "frame"); + const auto stream = at::cuda::getCurrentCUDAStream(); + + // === Step 1. The meshes of the frames, sized on the host === + std::vector frames(frame_count); + std::vector> plans(frame_count); + const double* cells = box.data_ptr(); + int64_t mesh_total = 0; + int64_t spectrum_size = 0; + int64_t factor_size = 0; + size_t work_bytes = 0; + for (int f = 0; f < frame_count; ++f) { + frames[f] = frame_mesh(cells + 9 * f, spacing); + frames[f].offset = mesh_total; + mesh_total += channels * mesh_points(frames[f]); + const int half = frames[f].n[2] / 2 + 1; + spectrum_size = std::max(spectrum_size, int64_t{channels} * frames[f].n[0] * + frames[f].n[1] * half); + factor_size = + std::max(factor_size, int64_t{frames[f].n[0]} + frames[f].n[1] + half); + plans[f] = mesh_plans(frames[f], channels); + work_bytes = std::max(work_bytes, plans[f]->work_bytes); + } + const auto byte_options = torch::TensorOptions().dtype(torch::kUInt8); + auto host_frames = + torch::empty({static_cast(sizeof(FrameMesh)) * frame_count}, + byte_options.pinned_memory(true)); + std::memcpy(host_frames.data_ptr(), frames.data(), + sizeof(FrameMesh) * frame_count); + const IntegerWeights integer_weight = + integer_bspline_weights(static_cast(order)); + // A few blocks per multiprocessor fill the device; every thread of the + // influence function then walks several modes, which keeps the block sums + // few. + const int mode_blocks = + 4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + + // === Step 2. The buffers of the call === + const int64_t average_nodes = (node_count + frame_count - 1) / frame_count; + const auto slices_for = [](int64_t nodes) { + return static_cast(std::min( + kMaxSlices, + std::max(1, (nodes + kSliceNodes - 1) / kSliceNodes))); + }; + const int slice_count = slices_for(average_nodes); + const int owner_slice_count = slices_for( + (int64_t{owner_count} + frame_count - 1) / std::max(frame_count, 1)); + const int64_t stencil_bytes = dispatch_order_value(order, [](auto constant) { + return static_cast(sizeof(AxisStencil)); + }); + Workspace workspace; + const int64_t frames_at = + workspace.reserve(sizeof(FrameMesh) * frame_count); + const int64_t frame_ptr_at = workspace.reserve(frame_count + 1); + const int64_t slice_charge_at = + workspace.reserve(int64_t{frame_count} * slice_count * 2); + const int64_t frame_stat_at = + workspace.reserve(int64_t{frame_count} * kStatWidth); + const int64_t stencils_at = + workspace.reserve(3 * node_count * stencil_bytes); + const int64_t node_charge_at = + workspace.reserve(node_count * channels); + const int64_t fixed_mesh_at = workspace.reserve(mesh_total); + // The potential mesh of a frame overwrites its charge mesh, which its + // forward transform has consumed. + const int64_t mesh_at = workspace.reserve(mesh_total); + const int64_t spectrum_at = workspace.reserve(spectrum_size); + const int64_t work_at = workspace.reserve(work_bytes); + const int64_t factor_at = workspace.reserve(factor_size); + const int64_t block_sum_at = + workspace.reserve(int64_t{mode_blocks} * kModeSums); + const int64_t frame_virial_at = + workspace.reserve(int64_t{frame_count} * 9); + const int64_t potential_at = + workspace.reserve(int64_t{owner_count} * channels); + const int64_t owner_slice_at = workspace.reserve( + int64_t{frame_count} * owner_slice_count * kOwnerSums); + const int64_t owner_local_at = + workspace.reserve(int64_t{frame_count} * kOwnerSums); + workspace.allocate(coord.device()); + auto* frames_ptr = workspace.at(frames_at); + auto* frame_ptr = workspace.at(frame_ptr_at); + auto* slice_charge = workspace.at(slice_charge_at); + auto* frame_stat = workspace.at(frame_stat_at); + auto* node_charge = workspace.at(node_charge_at); + auto* fixed_mesh = workspace.at(fixed_mesh_at); + auto* charge_mesh = workspace.at(mesh_at); + float* potential_mesh = charge_mesh; + auto* spectrum = workspace.at(spectrum_at); + void* work = workspace.at(work_at); + auto* factor = workspace.at(factor_at); + auto* block_sum = workspace.at(block_sum_at); + auto* frame_virial = workspace.at(frame_virial_at); + auto* potential = workspace.at(potential_at); + auto* owner_slice = workspace.at(owner_slice_at); + auto* owner_local = workspace.at(owner_local_at); + // The copy from the pinned staging buffer records its completion, before + // which the allocator does not hand the buffer out again. + torch::from_blob(frames_ptr, {host_frames.numel()}, + byte_options.device(coord.device())) + .copy_(host_frames, /*non_blocking=*/true); + C10_CUDA_CHECK( + cudaMemsetAsync(fixed_mesh, 0, sizeof(int64_t) * mesh_total, stream)); + + // === Step 3. The constraint, and the stencils and charges of the nodes === + const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); + frame_offsets<<<1, kOffsetThreads, 0, stream>>>(counts.data_ptr(), + frame_count, frame_ptr); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + charge_slices<<>>( + charge.data_ptr(), node_mask.data_ptr(), frame_ptr, channels, + static_cast(std::max(channel, 0)), slice_charge); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads, + kFrameThreads, 0, stream>>>( + slice_charge, slice_count, frame_count, total_charge.data_ptr(), + frame_stat); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + dispatch_order(order, [&](auto order_constant) { + constexpr int kOrder = decltype(order_constant)::value; + auto* node_stencil = workspace.at>(stencils_at); + node_stencils<<>>( + coord.data_ptr(), charge.data_ptr(), + node_mask.data_ptr(), frame_ptr, frames_ptr, frame_stat, + frame_count, node_count, channels, static_cast(channel), + node_stencil, node_charge); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + // === Step 4. The charge meshes === + spread_charges + <<>>( + node_mask.data_ptr(), frame_ptr, frames_ptr, node_stencil, + node_charge, frame_count, node_count, channels, fixed_mesh); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + fixed_to_float<<>>( + mesh_total, fixed_mesh, charge_mesh); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + // === Step 5. Influence function and potential meshes, frame by frame === + for (int f = 0; f < frame_count; ++f) { + const FrameMesh& grid = frames[f]; + MeshPlans& plan = *plans[f]; + { + const std::lock_guard lock(plan.mutex); + check_cufft(cufftSetStream(plan.forward, stream), "stream setup"); + check_cufft(cufftSetWorkArea(plan.forward, work), "work area setup"); + check_cufft(cufftExecR2C(plan.forward, charge_mesh + grid.offset, + reinterpret_cast(spectrum)), + "transform"); + } + const int half = grid.n[2] / 2 + 1; + const int factor_count = grid.n[0] + grid.n[1] + half; + assignment_factor<<>>( + grid, static_cast(order), integer_weight, factor); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + apply_influence<<>>( + grid, channels, static_cast(kernel), + coefficient.data_ptr(), gaussians, prefactor, factor, spectrum, + block_sum); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + mode_virial<<<1, kThreads, 0, stream>>>(block_sum, mode_blocks, + frame_virial + 9 * f); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + { + const std::lock_guard lock(plan.mutex); + check_cufft(cufftSetStream(plan.inverse, stream), "stream setup"); + check_cufft(cufftSetWorkArea(plan.inverse, work), "work area setup"); + check_cufft( + cufftExecC2R(plan.inverse, reinterpret_cast(spectrum), + potential_mesh + grid.offset), + "transform"); + } + } + + // === Step 6. Potential, force and virial of the owners === + auto force = torch::empty({owner_count, 3}, float_options); + if (owner_count > 0) { + dispatch_order(order, [&](auto order_constant) { + constexpr int kOrder = decltype(order_constant)::value; + interpolate_owners<<>>( + node_mask.data_ptr(), frame_ptr, frames_ptr, + workspace.at>(stencils_at), node_charge, + potential_mesh, frame_count, owner_begin, owner_count, channels, + potential, force.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // The frames of the owners, on the axis of the owners: the frames + // themselves when every node is an owner, else the one frame they share. + torch::Tensor owner_frame_tensor; + const int64_t* owner_frame_ptr = frame_ptr; + if (!every_node) { + owner_frame_tensor = + torch::full({2}, int64_t{owner_count}, counts.options()); + owner_frame_tensor.select(0, 0).zero_(); + owner_frame_ptr = owner_frame_tensor.data_ptr(); + } + const bool* owner_mask = node_mask.data_ptr() + owner_begin; + const double* owner_coord = coord.data_ptr() + 3 * owner_begin; + const float* owner_charge = node_charge + int64_t{owner_begin} * channels; + owner_slices<<>>(potential, force.data_ptr(), owner_coord, + owner_charge, owner_mask, owner_frame_ptr, channels, + static_cast(channel), static_cast(kernel), + coefficient.data_ptr(), gaussians, prefactor, + owner_slice); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + owner_totals<<>>(owner_slice, owner_slice_count, + owner_local); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + // The mean potential of the constraint and the virial share run over the + // charged nodes of every domain. + torch::Tensor owner_sum; + const double* owner_total = owner_local; + if (communicator != nullptr) { + owner_sum = sum_over_domains( + *communicator, torch::from_blob(owner_local, {frame_count * kOwnerSums}, + double_options)); + owner_total = owner_sum.data_ptr(); + } + auto atom_energy = torch::empty({owner_count, 1}, double_options); + auto charge_gradient = torch::empty({owner_count, channels}, float_options); + auto atom_virial = torch::empty({owner_count, 9}, float_options); + auto virial = torch::empty({frame_count, 9}, float_options); + finish_owners<<(owner_count, 9 * frame_count)), + kThreads, 0, stream>>>( + potential, force.data_ptr(), owner_coord, owner_charge, owner_mask, + owner_frame_ptr, frame_virial, owner_local, owner_total, frame_stat, + frame_count, owner_count, channels, static_cast(channel), + static_cast(kernel), coefficient.data_ptr(), gaussians, + prefactor, atom_energy.data_ptr(), + charge_gradient.data_ptr(), atom_virial.data_ptr(), + virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {atom_energy, charge_gradient, force, atom_virial, virial}; +} + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/ops.cc b/source/op/pt/long_range/ops.cc new file mode 100644 index 0000000000..364fa6c2c2 --- /dev/null +++ b/source/op/pt/long_range/ops.cc @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Operator schema of the latent-charge long-range energy. +// +// The schema is declared here, unconditionally, while the CUDA kernel +// registers itself only against a CUDA-enabled PyTorch (latent_charge_sum.cu), +// so that the Python front end can tell "library not loaded" from "this device +// has no kernel". +// +// For the nodes of the frames, the operator takes the raw charges c_i of the +// charge network in C float32 channels, the nodes that carry them (w_i = 1, +// and 0 elsewhere) and float64 positions r_i. The interactions carry the +// charges q_i = w_i c_i; when `channel` is a channel index rather than -1, +// that channel is shifted uniformly over the charged nodes of every frame so +// that it sums to the total charge Q_f of the frame, +// +// q_i^ch = w_i (c_i^ch - s_f), s_f = (sum_{j in f} w_j c_j^ch - Q_f) / +// max(n_f, 1), +// +// with n_f charged nodes in f. Without `box`, every frame is a complete +// non-periodic system, in which the potential of the charges is the direct +// sum +// +// phi_i = prefactor * sum_{j != i} q_j K(r_ij); +// +// with `box`, whose row f holds the row-major cell of frame f, the rows of the +// cell being its lattice vectors, phi_i sums over the images of the cell, by +// the smooth particle-mesh Ewald method (mesh_sum.cu), less the potential +// prefactor * K(0) q_i of every charge at itself. An all-zero row marks a +// non-periodic frame, which a call does not mix with periodic ones. Either way +// the operator returns +// +// atom_energy_i = q_i . phi_i / 2 (float64) +// charge_gradient_i = dE / dc_i = w_i (phi_i - [c = ch] m_f), m_f the +// mean of phi^ch over the charged nodes of f +// force_i = -dE / dr_i at fixed charges +// atom_virial_i (9 components, row-major) +// virial_f = sum_{i in f} atom_virial_i +// +// in float32 unless noted. In a non-periodic frame, with d_ij = r_i - r_j, +// +// force_i = -prefactor * sum_{j != i} (q_i . q_j) K'(r_ij) / r_ij * d_ij, +// W_i = -prefactor / 2 * sum_{j != i} (q_i . q_j) K'(r_ij) / r_ij +// * d_ij (x) d_ij; +// +// in a periodic frame the virial of the frame is that of the strain of the +// cell, W_f = -dE / d(strain), and W_i = F_i (x) r_i + (W_f - sum_{j in f} +// F_j (x) r_j) / n_f on the charged nodes, the force index first, so that the +// charged nodes of the frame sum to W_f. +// +// The frames occupy the leading sum(n_node_per_frame) <= N nodes, which the +// caller ensures; the nodes past them are padding and join none. The frame +// energy is the sum of the atomic energies. +// +// `kernel` selects K: 0 is erf(alpha r) / r with `coefficient` holding alpha; +// 1 is a sum of Gaussians a_k 2^(-c_k r^2), c_k = log2(e) / (2 sigma_k^2), +// with `coefficient` holding the amplitudes a_k and then the rates c_k. +// `coefficient` is float64. `total_charge` holds Q_f in float64 and is read +// only with a constrained channel. `box` is a float64 CPU tensor of shape +// (F, 9). The mesh of a periodic frame holds along every lattice vector a_d +// the smallest even number of points without a prime factor above 7 of at +// least ceil(|a_d| / spacing), and spreads the charges with B-splines of +// order `order`. The force and the virial hold the charges fixed; the response +// of the charges is the caller's, through the charge gradient. +// +// `communicator` is an int64 CPU tensor of one element holding the address of +// the caller's MPI_Comm handle, or zero. A handle of several ranks decomposes +// one frame into domains, one per rank: the charged nodes of every rank are +// the charged nodes of the frame, and the sums run over all of them. The +// outputs of a rank are those of its charged nodes, and W_f is their share of +// the virial of the frame; the other nodes of the rank carry no charge. + +#include + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.def( + "latent_charge_sum(Tensor coord, Tensor charge, Tensor node_mask, " + "Tensor n_node_per_frame, Tensor total_charge, Tensor? box, " + "Tensor coefficient, Tensor communicator, int kernel, int channel, " + "float prefactor, int order, float spacing) -> " + "(Tensor atom_energy, Tensor charge_gradient, Tensor force, " + "Tensor atom_virial, Tensor virial)"); +} + +// The operator serves inference, whose caller differentiates nothing through +// it; its outputs carry no gradient. +TORCH_LIBRARY_IMPL(deepmd, Autograd, library) { + library.impl("latent_charge_sum", torch::CppFunction::makeFallthrough()); +} diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py index 5630fd5c61..2733c05346 100644 --- a/source/tests/common/dpmodel/test_linear_atomic_model.py +++ b/source/tests/common/dpmodel/test_linear_atomic_model.py @@ -47,6 +47,9 @@ def get_dim_fparam(self) -> int: def get_dim_aparam(self) -> int: return 0 + def has_long_range(self) -> bool: + return False + def change_type_map( self, type_map: list[str], diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py new file mode 100644 index 0000000000..2ba5dffb62 --- /dev/null +++ b/source/tests/common/dpmodel/test_long_range.py @@ -0,0 +1,918 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Primitives of the long-range modules: pair topology, kernels and constraint.""" + +import itertools + +import numpy as np +import pytest +from scipy.special import ( + erf, +) + +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting, +) +from deepmd.dpmodel.long_range import ( + COULOMB_CONSTANT, + BaseInteractionKernel, + LatentChargeLongRange, + LESKernel, + SOGKernel, +) +from deepmd.dpmodel.model.model import ( + get_model, +) +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + K_BLOCK, + build_k_grid, + build_pair_graph, + build_topologies, + enumerate_k_grid_bounds, + pair_force_virial, + unordered_pair_index, +) +from deepmd.utils.argcheck import ( + long_range_args, +) + + +def _brute_force_pairs(counts: list[int], mask: np.ndarray | None) -> np.ndarray: + pairs = [] + start = 0 + for count in counts: + nodes = [ + node for node in range(start, start + count) if mask is None or mask[node] + ] + pairs += itertools.combinations(nodes, 2) + start += count + return np.asarray(pairs, dtype=np.int64).reshape(-1, 2).T + + +@pytest.mark.parametrize( + "counts", [[0], [1], [2], [3, 0, 1, 5], [7, 200, 1, 0, 33], [1000]] +) +@pytest.mark.parametrize("masked", [False, True]) +def test_pair_index_matches_brute_force(counts: list[int], masked: bool) -> None: + """Every unordered pair of valid nodes of a frame appears once, in row-major order.""" + mask = np.random.default_rng(0).random(sum(counts)) > 0.3 if masked else None + index = unordered_pair_index(np.asarray(counts), mask) + np.testing.assert_array_equal(index, _brute_force_pairs(counts, mask)) + + +def test_pair_graph_guards_and_periodic_rejection() -> None: + """The pair axis ends in masked guards; a cell is rejected.""" + coord = np.random.default_rng(1).normal(size=(4, 3)) + atype = np.array([0, 1, 1, 0]) + pairs = build_pair_graph(coord, np.array([4]), atype) + assert pairs.index.shape == (2, 8) + assert pairs.mask.tolist() == [True] * 6 + [False] * 2 + np.testing.assert_allclose( + pairs.vec[:6], coord[pairs.index[0, :6]] - coord[pairs.index[1, :6]] + ) + assert build_topologies((), coord, np.array([4]), atype=atype) == {} + with pytest.raises(NotImplementedError, match="non-periodic"): + build_topologies( + ("all_pairs",), coord, np.array([4]), np.eye(3)[None], atype=atype + ) + + +def test_pair_force_virial_of_a_central_potential() -> None: + """For E = sum 1/r the assembled force is Newtonian and the virial sums.""" + coord = np.array([[0.0, 0.0, 0.0], [1.0, 0.2, 0.0], [0.0, 1.5, 0.3]]) + pairs = build_pair_graph(coord, np.array([3]), np.array([0, 1, 0])) + # d(1/r)/d(vec) on the real pairs; the zero-length guards carry no gradient. + r = np.where(pairs.mask, np.linalg.norm(pairs.vec, axis=-1), 1.0) + grad = np.where(pairs.mask[:, None], -pairs.vec / r[:, None] ** 3, 0.0) + force, atom_virial, virial = pair_force_virial( + grad, pairs.vec, pairs.index, pairs.mask, np.array([3]) + ) + expected = np.zeros((3, 3)) + for i, j in itertools.permutations(range(3), 2): + d = coord[i] - coord[j] + expected[i] += d / np.linalg.norm(d) ** 3 + np.testing.assert_allclose(force, expected, rtol=1e-12) + np.testing.assert_allclose(atom_virial.sum(0), virial[0], rtol=1e-12) + + +def test_les_kernel_closed_form() -> None: + """LES is erf(alpha r)/r, finite at contact.""" + kernel = LESKernel(alpha=0.7) + r = np.array([0.0, 0.5, 2.0, 9.0]) + values = kernel(r**2) + np.testing.assert_allclose(values[1:], erf(0.7 * r[1:]) / r[1:], rtol=1e-10) + np.testing.assert_allclose(values[0], 2 * 0.7 / np.sqrt(np.pi), rtol=1e-10) + + +def test_sog_kernel_closed_form_and_defaults() -> None: + """SOG is sum_k a_k exp(-r^2 / 2 sigma_k^2); the defaults discretize 1/r.""" + kernel = SOGKernel(amplitude=[0.7, -0.2], width=[0.8, 1.7]) + r = np.array([0.0, 1.1, 3.0]) + expected = 0.7 * np.exp(-(r**2) / (2 * 0.8**2)) - 0.2 * np.exp( + -(r**2) / (2 * 1.7**2) + ) + np.testing.assert_allclose(kernel(r**2), expected, rtol=1e-12) + default = SOGKernel(n_gaussian=3, base_width=2.0, width_ratio=1.5) + np.testing.assert_allclose(np.exp(default.log_width), [2.0, 3.0, 4.5]) + np.testing.assert_allclose( + default.amplitude, np.sqrt(2 / np.pi) * np.log(1.5) / np.array([2.0, 3.0, 4.5]) + ) + # The default kernel exerts the Coulomb force, r^2 K'(r) = -1, between a + # few times its first width and a fraction of its widest one. + default = SOGKernel() + r = np.geomspace(6.5, 150.0, 200) + step = 1e-4 * r + slope = (default((r + step) ** 2) - default((r - step) ** 2)) / (2 * step) + np.testing.assert_allclose(r**2 * slope, -1.0, rtol=0, atol=5e-3) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"type": "les", "alpha": 0.0}, + {"type": "sog", "width": [1.0, -2.0]}, + {"type": "sog", "width": [1.0, 2.0], "amplitude": [1.0, 2.0, 3.0]}, + {"type": "sog", "width": [1.0, 2.0]}, + {"type": "sog", "width_ratio": 1.0}, + {"type": "sog", "amplitude": 1.0, "width_ratio": 0.5}, + ], +) +def test_kernel_parameter_domain(kwargs: dict) -> None: + """Widths must be positive and increase by default; amplitudes match the Gaussians.""" + kernel_cls = BaseInteractionKernel.get_class_by_type(kwargs.pop("type")) + with pytest.raises(ValueError): + kernel_cls(**kwargs) + + +@pytest.mark.parametrize("kernel", [LESKernel(alpha=1.3), SOGKernel(n_gaussian=4)]) +def test_kernel_round_trip(kernel: BaseInteractionKernel) -> None: + """A kernel record restores the same parameters.""" + restored = BaseInteractionKernel.deserialize(kernel.serialize()) + r_sq = np.linspace(0.0, 30.0, 7) + np.testing.assert_allclose(restored(r_sq), kernel(r_sq), rtol=0, atol=0) + + +def test_kernel_defaults_to_les() -> None: + """A configuration without a kernel, or without its type, selects LES.""" + args = long_range_args() + config = args.normalize_value({"type": "latent_charge"}, trim_pattern="_*") + assert config["kernel"] == {"type": "les", "alpha": 1.0, "trainable": True} + config = args.normalize_value( + {"type": "latent_charge", "kernel": {"alpha": 2.0}}, trim_pattern="_*" + ) + assert config["kernel"]["type"] == "les" + model = get_model( + { + "type_map": ["A", "B"], + "descriptor": {"type": "dpa4c", "rcut": 3.0, "channels": 8, "seed": 1}, + "fitting_net": {"type": "ener", "neuron": [8], "seed": 2}, + "long_range": {"type": "latent_charge", "charge_net": {"neuron": [4]}}, + } + ) + kernel = model.atomic_model.long_range.kernel + assert isinstance(kernel, LESKernel) + np.testing.assert_allclose(np.exp(kernel.log_alpha), 1.0) + + +def _long_range(constraint: dict | None, channels: int = 2) -> LatentChargeLongRange: + charge_net = LatentChargeFitting( + "latent_charge", 2, 4, channels, neuron=[4], numb_fparam=1, seed=1 + ) + return LatentChargeLongRange(charge_net, LESKernel(), constraint) + + +def test_constraint_shifts_one_channel_over_the_charged_atoms() -> None: + """The constrained channel sums to the target over the charged atoms only.""" + long_range = _long_range({"source": "fparam", "channel": 1, "fparam_index": 0}) + rng = np.random.default_rng(2) + charge = rng.normal(size=(7, 2)) + node_mask = np.array([True, True, False, True, True, True, False]) + charge[~node_mask] = 0.0 + n_node = np.array([4, 3]) + frame_id = np.array([0, 0, 0, 0, 1, 1, 1]) + fparam = np.array([[1.5], [-2.0]]) + constrained = long_range._constrain( + charge, node_mask, frame_id, n_node, fparam, None + ) + np.testing.assert_array_equal(constrained[:, 0], charge[:, 0]) + np.testing.assert_array_equal(constrained[~node_mask], 0.0) + totals = np.bincount(frame_id, weights=constrained[:, 1], minlength=2) + np.testing.assert_allclose(totals, [1.5, -2.0], rtol=1e-12) + # A frame without charged atoms is left as it is. + empty = long_range._constrain( + charge * 0.0, + np.zeros(7, dtype=bool), + frame_id, + n_node, + fparam, + None, + ) + np.testing.assert_array_equal(empty, 0.0) + + +@pytest.mark.parametrize( + "constraint", + [ + {"source": "total"}, + {"source": "neutral", "channel": 2}, + {"source": "fparam", "fparam_index": 1}, + ], +) +def test_constraint_configuration_is_validated(constraint: dict) -> None: + """Unknown sources and out-of-range channels or columns are rejected.""" + with pytest.raises(ValueError): + _long_range(constraint) + + +def test_long_range_round_trip() -> None: + """The module record holds the charge network, kernel and constraint.""" + long_range = _long_range({"source": "neutral"}) + record = long_range.serialize() + restored = LatentChargeLongRange.deserialize(record) + assert restored.serialize()["charge_constraint"] == record["charge_constraint"] + assert restored.required_topologies() == ("all_pairs",) + assert not restored.uses_charge_spin() + + +# --------------------------------------------------------------------------- +# Periodic (k-space) primitives +# --------------------------------------------------------------------------- + + +def _spherical_ft( + kernel: str, k: float, n_r: int = 20000, r_max: float = 60.0 +) -> float: + """Isotropic Fourier transform of the kernel by direct radial quadrature.""" + r = np.linspace(1e-9, r_max, n_r) + if kernel == "les": + k_r = erf(0.8 * r) / r + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k_r = np.sum( + amplitude[None, :] * np.exp(-(r[:, None] ** 2) / (2 * width**2)), -1 + ) + return float(np.trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(k * r / np.pi), r)) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_kernel_kspace_matches_radial_ft(kernel: str) -> None: + """``kernel.kspace`` is the analytic Fourier transform of the real-space kernel.""" + k_sq = np.array([0.5, 1.0, 2.0, 5.0]) ** 2 + if kernel == "les": + mod = LESKernel(alpha=0.8) + else: + mod = SOGKernel(amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2]) + got = mod.kspace(k_sq) + for kk, value in zip(np.sqrt(k_sq), got, strict=True): + if kernel == "les": + # erf(alpha r)/r is not integrable; quadrature of the integrable + # complement erfc(alpha r)/r = 1/r - erf(alpha r)/r converges, and + # K_hat(k) = 4 pi / k^2 - FT(erfc(alpha r)/r). + r = np.linspace(1e-9, 60.0, 20000) + k_r = (1.0 - erf(0.8 * r)) / r # erfc(0.8 r)/r + expected = 4.0 * np.pi / kk**2 - float( + np.trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(kk * r / np.pi), r) + ) + else: + expected = _spherical_ft(kernel, kk) + np.testing.assert_allclose(value, expected, rtol=1e-3, atol=1e-5) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_kernel_contact_value_matches_real_space_limit(kernel: str) -> None: + """``kernel.contact_value`` is K(r) at contact.""" + if kernel == "les": + mod = LESKernel(alpha=0.8) + else: + mod = SOGKernel(amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2]) + r_sq = np.array([1e-12]) + np.testing.assert_allclose( + mod.contact_value(r_sq), mod(r_sq)[0], rtol=1e-6, atol=1e-9 + ) + + +def test_build_k_grid_orthogonality_and_guards() -> None: + """The k points satisfy k . a_d = 2 pi n_d, every frame fills a row of its own.""" + box = np.array([[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]])[None] + grid = build_k_grid(box, 2.0, np.zeros((4, 3)), np.array([4])) + # the points fill whole blocks, the last one completed by guard points + assert grid.mask.shape == grid.vec.shape[:3] + assert grid.vec.shape[2] == K_BLOCK + flat_mask = grid.mask.reshape(-1) + count = int(flat_mask.sum()) + assert flat_mask[:count].all() and not flat_mask[count:].any() + np.testing.assert_array_equal(grid.slot, [[0, 1, 2, 3]]) + vec = grid.vec.reshape(-1, 3)[:count] + n_integer = vec @ box[0].T / (2 * np.pi) + np.testing.assert_allclose(n_integer, np.round(n_integer), atol=1e-10) + k_sq = np.sum(vec**2, axis=-1) + assert (k_sq > 0).all() and (k_sq <= 4.0 + 1e-12).all() + np.testing.assert_allclose(grid.volume, np.abs(np.linalg.det(box[0]))) + # every integer point inside the sphere is present exactly once + ns = np.round(n_integer).astype(int) + assert len({tuple(row) for row in ns}) == ns.shape[0] + rec = 2 * np.pi * np.linalg.inv(box[0]) + n_d = np.ceil(2.0 * np.linalg.norm(box[0], axis=1) / (2 * np.pi)).astype(int) + count = 0 + for n0 in range(-n_d[0], n_d[0] + 1): + for n1 in range(-n_d[1], n_d[1] + 1): + for n2 in range(-n_d[2], n_d[2] + 1): + if (n0, n1, n2) != (0, 0, 0): + k = rec @ np.array([n0, n1, n2], dtype=float) + count += int(np.sum(k * k) <= 4.0) + assert count == ns.shape[0] + + +def test_build_topologies_periodic_branching() -> None: + """A cell switches the pair topology to guards and the k grid on, and back.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + K_GRID, + ) + + coord = np.random.default_rng(1).normal(size=(4, 3)) + atype = np.array([0, 1, 1, 0]) + n_node = np.array([4]) + box = np.eye(3)[None] * 10.0 + names = (ALL_PAIRS, K_GRID) + params = {K_GRID: {"k_max": 2.0}} + with_box = build_topologies( + names, coord, n_node, box, atype=atype, topology_params=params + ) + assert with_box[ALL_PAIRS].index.shape[1] == 4 # guards only + assert not with_box[ALL_PAIRS].mask.any() + assert with_box[K_GRID].mask.sum() > 0 + without_box = build_topologies( + names, coord, n_node, None, atype=atype, topology_params=params + ) + assert without_box[ALL_PAIRS].mask.any() + assert not without_box[K_GRID].mask.any() # guards only + # a pair-only declaration still rejects a cell loudly + with pytest.raises(NotImplementedError, match="non-periodic"): + build_topologies((ALL_PAIRS,), coord, n_node, box, atype=atype) + # the k grid needs its build parameter + with pytest.raises(ValueError, match="k_max"): + build_topologies(names, coord, n_node, box, atype=atype) + # the reciprocal-space sum excludes no pair; a non-periodic batch still does + pair_excl = PairExcludeMask(2, [(0, 1)]) + with pytest.raises(NotImplementedError, match="pair_exclude_types"): + build_topologies( + names, + coord, + n_node, + box, + atype=atype, + pair_excl=pair_excl, + topology_params=params, + ) + excluded = build_topologies( + names, coord, n_node, None, atype=atype, pair_excl=pair_excl + ) + assert 0 < excluded[ALL_PAIRS].mask.sum() < without_box[ALL_PAIRS].mask.sum() + # a cell without volume is no periodic system + with pytest.raises(ValueError, match="no volume"): + build_topologies( + names, coord, n_node, np.zeros((1, 9)), atype=atype, topology_params=params + ) + + +def test_build_k_grid_pads_the_rows_of_the_frames() -> None: + """Every frame keeps its own points and nodes; shorter rows end in masked guards.""" + box = np.stack([np.diag([8.0, 7.5, 9.0]), np.diag([5.0, 5.5, 6.0])]) + n_node = np.array([3, 5]) + grid = build_k_grid(box, 2.0, np.zeros((8, 3)), n_node) + alone = [ + build_k_grid(box[f : f + 1], 2.0, np.zeros((n, 3)), np.array([n])) + for f, n in enumerate(n_node) + ] + counts = [int(single.mask.sum()) for single in alone] + blocks = max(-(-max(counts) // K_BLOCK), 2) + assert grid.vec.shape == (2, blocks, K_BLOCK, 3) + for f, single in enumerate(alone): + mask = grid.mask[f].reshape(-1) + vec = grid.vec[f].reshape(-1, 3) + np.testing.assert_array_equal(mask, np.arange(blocks * K_BLOCK) < counts[f]) + np.testing.assert_allclose( + vec[: counts[f]], single.vec[0].reshape(-1, 3)[: counts[f]] + ) + assert np.all(vec[counts[f] :] == [1e-6, 0.0, 0.0]) + np.testing.assert_array_equal(grid.slot, [[0, 1, 2, -1, -1], [3, 4, 5, 6, 7]]) + + +@pytest.mark.parametrize("kernel", [LESKernel(alpha=0.6), SOGKernel(n_gaussian=3)]) +def test_kspace_energy_of_padded_frames_matches_frames_alone( + kernel: BaseInteractionKernel, +) -> None: + """Frames of unequal node counts, an empty one among them, sum as they do alone.""" + from deepmd.dpmodel.long_range.reciprocal import ( + kspace_energy, + ) + + box = np.stack( + [np.diag([8.0, 7.5, 9.0]), np.diag([6.0, 6.5, 7.0]), np.diag([5.0, 5.5, 6.0])] + ) + n_node = np.array([4, 0, 6]) + rng = np.random.default_rng(3) + coord = rng.uniform(0.0, 5.0, (10, 3)) + charge = rng.normal(0.0, 0.5, (10, 2)) + energy = kspace_energy(charge, build_k_grid(box, 2.5, coord, n_node), kernel) + start = 0 + for f, n in enumerate(n_node): + if n == 0: + continue + nodes = slice(start, start + n) + alone = build_k_grid(box[f : f + 1], 2.5, coord[nodes], np.array([n])) + np.testing.assert_allclose( + energy[nodes], kspace_energy(charge[nodes], alone, kernel), rtol=1e-12 + ) + start += n + + +# --------------------------------------------------------------------------- +# Fixed k-grid mode +# --------------------------------------------------------------------------- + + +def _module_topologies( + module: LatentChargeLongRange, + coord: np.ndarray, + n_node: np.ndarray, + box: np.ndarray | None, + atype: np.ndarray, +) -> dict: + """The topologies of a module, built with its parameters.""" + return build_topologies( + module.required_topologies(), + coord, + n_node, + box, + atype=atype, + topology_params=module.topology_params(), + ) + + +def _periodic_long_range(**kwargs: object) -> LatentChargeLongRange: + """A periodic module, over the explicit wave numbers unless ``k_grid_mode`` says otherwise.""" + charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1) + kwargs.setdefault("k_grid_mode", "ragged") + return LatentChargeLongRange( + charge_net, LESKernel(alpha=0.3), None, periodic=True, **kwargs + ) + + +def _cells() -> np.ndarray: + """Two cells of different sizes, the second the larger one.""" + return np.array( + [ + [[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]], + [[12.0, 0.5, 0.0], [0.2, 11.0, 0.3], [0.1, 0.4, 10.5]], + ] + ) + + +def test_enumerate_k_grid_bounds() -> None: + """The grid is the lexicographic meshgrid of the bounds, without the zero vector.""" + n_grid = enumerate_k_grid_bounds([2, 1, 0]) + assert n_grid.shape == (5 * 3 * 1 - 1, 3) + assert not np.any(np.all(n_grid == 0, axis=-1)) + expected = [ + (n0, n1, 0) for n0 in range(-2, 3) for n1 in range(-1, 2) if (n0, n1) != (0, 0) + ] + np.testing.assert_array_equal(n_grid, np.asarray(expected, dtype=np.int64)) + with pytest.raises(ValueError, match="bounds"): + enumerate_k_grid_bounds([1, -1, 0]) + + +def test_fixed_k_grid_matches_ragged_enumeration() -> None: + """The masked points of the fixed grid are exactly the ragged k points.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_fixed_k_grid, + enumerate_k_grid_bounds, + k_grid_bounds_of_cells, + ) + + box = _cells() + k_max = 2.0 + bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + nk = enumerate_k_grid_bounds(bounds).shape[0] + n_node = np.array([4, 5]) + fixed = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3)), n_node) + # a static layout: the blocks of the grid in every row + assert fixed.vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3) + ragged = build_k_grid(box, k_max, np.zeros((9, 3)), n_node) + for frame in range(2): + expected = { + tuple(row) for row in np.round(ragged.vec[frame][ragged.mask[frame]], 8) + } + got = {tuple(row) for row in np.round(fixed.vec[frame][fixed.mask[frame]], 8)} + assert got == expected + # the smaller cell masks more points of the shared grid + assert fixed.mask[0].sum() < fixed.mask[1].sum() + np.testing.assert_allclose(fixed.volume, np.abs(np.linalg.det(box)), rtol=1e-12) + + +def test_fixed_k_grid_truncates_oversized_cells() -> None: + """A cell beyond the grid is served by the grid alone (truncated k-sum), with a warning.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_fixed_k_grid, + enumerate_k_grid_bounds, + k_grid_bounds_of_cells, + ) + + box = _cells() + k_max = 2.0 + bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + nk = enumerate_k_grid_bounds(bounds).shape[0] + n_node = np.array([4, 5]) + with pytest.warns(UserWarning, match="truncated"): + bigger = build_fixed_k_grid(box * 2.0, bounds, k_max, np.zeros((9, 3)), n_node) + # the layout is untouched: static, every grid point present + assert bigger.vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3) + # every grid point is still a valid wave number of the larger cell; the + # mask keeps the whole k_max sphere the grid reaches + k2 = np.sum(bigger.vec**2, axis=-1) + real = (np.arange(bigger.mask[0].size) < nk).reshape(bigger.mask.shape[1:]) + np.testing.assert_array_equal(bigger.mask, real & (k2 <= k_max * k_max)) + # the larger cell activates more of the shared grid than the original did + original = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3)), n_node) + assert bigger.mask.sum() > original.mask.sum() + # a larger wave-number cutoff (a drifting kernel) truncates the same way + with pytest.warns(UserWarning, match="truncated"): + build_fixed_k_grid(box, bounds, k_max * 2.0, np.zeros((9, 3)), n_node) + + +def test_fixed_k_grid_mode_configuration() -> None: + """The fixed mode requires periodicity, and the bounds the fixed mode.""" + with pytest.raises(ValueError, match="k_grid_mode"): + _periodic_long_range(k_grid_mode="dense") + with pytest.raises(ValueError, match="periodic"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + k_grid_mode="fixed", + ) + with pytest.raises(ValueError, match="k_grid_bounds"): + _periodic_long_range(k_grid_bounds=[3, 3, 3]) + # explicit bounds build the grid at construction, no statistics needed + long_range = _periodic_long_range(k_grid_mode="fixed", k_grid_bounds=[2, 2, 2]) + np.testing.assert_array_equal(long_range.k_grid_bounds, [2, 2, 2]) + with pytest.raises(ValueError, match="non-negative"): + _periodic_long_range(k_grid_mode="fixed", k_grid_bounds=[0, 0, 0]) + + +def test_fixed_k_grid_stats_size_the_grid() -> None: + """The data statistics size the grid by the largest sampled cell.""" + long_range = _periodic_long_range(k_grid_mode="fixed") + # the grid is not sized before the statistics run + coord = np.zeros((9, 3)) + atype = np.zeros(9, dtype=np.int64) + with pytest.raises(ValueError, match="not sized yet"): + _module_topologies(long_range, coord, np.array([4, 5]), _cells(), atype) + long_range.compute_input_stats( + [{"box": _cells()[:1].reshape(1, 9)}, {"box": _cells()}] + ) + k_max = long_range.kernel.kspace_cutoff(long_range.kspace_tolerance) + from deepmd.dpmodel.utils.neighbor_graph import ( + k_grid_bounds_of_cells, + ) + + expected = np.max(k_grid_bounds_of_cells(_cells(), k_max), axis=0) + np.testing.assert_array_equal(long_range.k_grid_bounds, expected) + topologies = _module_topologies( + long_range, coord, np.array([4, 5]), _cells(), atype + ) + nk = np.prod(2 * expected + 1) - 1 + assert topologies["k_grid"].vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3) + # statistics without a cell cannot size the grid + unsized = _periodic_long_range(k_grid_mode="fixed") + with pytest.raises(ValueError, match="no simulation cell"): + unsized.compute_input_stats([{"coord": np.zeros((1, 3))}]) + + +def test_fixed_k_grid_round_trip() -> None: + """The record keeps the grid mode and the sized grid.""" + long_range = _periodic_long_range(k_grid_mode="fixed") + long_range.compute_input_stats([{"box": _cells()}]) + record = long_range.serialize() + assert record["k_grid_mode"] == "fixed" + restored = LatentChargeLongRange.deserialize(record) + assert restored.k_grid_mode == "fixed" + np.testing.assert_array_equal(restored.k_grid_bounds, long_range.k_grid_bounds) + # the ragged mode serializes no grid bounds + plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize()) + assert plain.k_grid_mode == "ragged" + assert plain.k_grid_bounds is None + # a periodic module evaluates on the mesh unless told otherwise, and a + # periodic record without a mode is one of the explicit wave numbers + charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1) + default = LatentChargeLongRange(charge_net, LESKernel(), None, periodic=True) + assert default.k_grid_mode == "pme" + assert default.required_topologies() == ("all_pairs", "pme_mesh") + record = default.serialize() + assert record["k_grid_mode"] == "pme" + record.pop("k_grid_mode") + assert LatentChargeLongRange.deserialize(record).k_grid_mode == "ragged" + + +# --------------------------------------------------------------------------- +# Decoupled k grid +# --------------------------------------------------------------------------- + + +def test_decoupled_k_grid_requires_periodic() -> None: + """Decoupling is only meaningful for a periodic module.""" + with pytest.raises(ValueError, match="decouple_k_grid"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + decouple_k_grid=True, + ) + + +def test_decoupled_k_grid_ignores_kernel_drift() -> None: + """A decoupled grid keeps the k points of the initial kernel however the + kernel parameters move; a coupled one re-derives them every batch. + """ + coord = np.random.default_rng(3).normal(size=(5, 3)) + atype = np.zeros(5, dtype=np.int64) + n_node = np.array([5]) + box = np.eye(3)[None] * 8.0 + + decoupled = _periodic_long_range(decouple_k_grid=True) + coupled = _periodic_long_range() + # the frozen cutoff is the cutoff of the initial kernel + assert decoupled.frozen_k_max == pytest.approx( + LESKernel(alpha=0.3).kspace_cutoff(1e-8) + ) + initial = _module_topologies(decoupled, coord, n_node, box, atype)["k_grid"] + reference = _module_topologies(coupled, coord, n_node, box, atype)["k_grid"] + np.testing.assert_array_equal(initial.vec, reference.vec) + # a larger alpha moves the cutoff out, growing the coupled grid + for module in (decoupled, coupled): + module.kernel.log_alpha = np.log(1.5) + drifted = _module_topologies(decoupled, coord, n_node, box, atype)["k_grid"] + moved = _module_topologies(coupled, coord, n_node, box, atype)["k_grid"] + np.testing.assert_array_equal(drifted.vec, initial.vec) + np.testing.assert_array_equal(drifted.mask, initial.mask) + assert np.sum(moved.mask) > np.sum(initial.mask) + + +def test_decoupled_fixed_grid_keeps_its_mask() -> None: + """The mask of a decoupled fixed grid ignores the kernel drift.""" + coord = np.zeros((5, 3)) + atype = np.zeros(5, dtype=np.int64) + n_node = np.array([5]) + box = np.eye(3)[None] * 8.0 + + long_range = _periodic_long_range( + k_grid_mode="fixed", k_grid_bounds=[5, 5, 5], decouple_k_grid=True + ) + before = _module_topologies(long_range, coord, n_node, box, atype)["k_grid"] + # a smaller alpha moves the cutoff in; a coupled mask would tighten + long_range.kernel.log_alpha = np.log(0.05) + after = _module_topologies(long_range, coord, n_node, box, atype)["k_grid"] + np.testing.assert_array_equal(after.mask, before.mask) + np.testing.assert_array_equal(after.vec, before.vec) + + +def test_decoupled_k_grid_round_trip() -> None: + """The frozen cutoff survives serialization, past the kernel drift.""" + long_range = _periodic_long_range(decouple_k_grid=True) + frozen = long_range.frozen_k_max + long_range.kernel.log_alpha = np.log(1.5) + record = long_range.serialize() + assert record["decouple_k_grid"] is True + assert record["frozen_k_max"] == pytest.approx(frozen) + restored = LatentChargeLongRange.deserialize(record) + assert restored.decouple_k_grid + assert restored.frozen_k_max == pytest.approx(frozen) + # ...even though the restored (drifted) kernel would derive a different one + assert restored.kernel.kspace_cutoff(1e-8) != pytest.approx(frozen) + # the default stays coupled and serializes nothing + plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize()) + assert not plain.decouple_k_grid + assert plain.frozen_k_max is None + + +# --------------------------------------------------------------------------- +# PME mesh mode +# --------------------------------------------------------------------------- + + +def _pme_long_range(**kwargs: object) -> LatentChargeLongRange: + kwargs.setdefault("k_grid_mode", "pme") + return _periodic_long_range(**kwargs) + + +def test_pme_mesh_grid_build_and_flatten() -> None: + """The mesh grid packs the cells and the mesh, flags the frames and flattens.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, + build_topologies, + flatten_topologies, + mesh_dims_of_cells, + unflatten_topologies, + ) + + coord = np.random.default_rng(1).normal(size=(4, 3)) + atype = np.array([0, 1, 1, 0]) + box = _cells() + params = {PME_MESH: {"spacing": 1.0}} + tops = build_topologies( + ("all_pairs", PME_MESH), + coord, + np.array([2, 2]), + box, + atype=atype, + topology_params=params, + ) + mesh = tops[PME_MESH] + np.testing.assert_array_equal(mesh.box, box) + assert mesh.mask.shape == (2,) and mesh.mask.all() + dims = tuple(m.shape[0] for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)) + assert dims == mesh_dims_of_cells(box, 1.0) + # the modes of an axis are its integer frequencies in FFT order + n1 = dims[0] + expected = np.arange(n1, dtype=float) + np.testing.assert_array_equal( + mesh.mode_1, np.where(expected > n1 // 2, expected - n1, expected) + ) + assert tops["all_pairs"].index.shape[1] == 4 # guards only + flat = flatten_topologies(tops) + assert len(flat) == 3 + 8 + back = unflatten_topologies(("all_pairs", PME_MESH), flat) + np.testing.assert_array_equal(back[PME_MESH].box, box) + np.testing.assert_array_equal(back[PME_MESH].mask, mesh.mask) + np.testing.assert_array_equal(back[PME_MESH].mode_3, mesh.mode_3) + # a non-periodic batch: real pairs, masked-out mesh + tops_np = build_topologies( + ("all_pairs", PME_MESH), + coord, + np.array([2, 2]), + None, + atype=atype, + topology_params=params, + ) + assert tops_np["all_pairs"].mask.any() + assert not tops_np[PME_MESH].mask.any() + # missing build parameters are an error + with pytest.raises(ValueError, match="spacing"): + build_topologies((PME_MESH,), coord, np.array([4]), box, atype=atype) + + +def test_pme_mesh_dims_resolve_every_cell() -> None: + """The mesh is the smallest even 7-smooth one resolving every cell.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + mesh_dims_of_cells, + ) + + for spacing in (0.37, 0.5, 1.1): + dims = mesh_dims_of_cells(_cells(), spacing) + longest = np.max(np.linalg.norm(_cells(), axis=-1), axis=0) + for n, a in zip(dims, longest, strict=True): + assert n % 2 == 0 and n >= a / spacing + m = n + for p in (2, 3, 5, 7): + while m % p == 0: + m //= p + assert m == 1 + # no smaller even 7-smooth size reaches the spacing + for smaller in range(n - 2, 1, -2): + m = smaller + for p in (2, 3, 5, 7): + while m % p == 0: + m //= p + if m == 1: + assert smaller < a / spacing + break + # the larger cell alone sizes the mesh of the batch + assert mesh_dims_of_cells(_cells(), 0.5) == mesh_dims_of_cells(_cells()[1:], 0.5) + + +def _pme_energy(module, charge, coord, box) -> float: + from deepmd.dpmodel.long_range.reciprocal import ( + mesh_energy, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_mesh_grid, + ) + + mesh = build_mesh_grid(box[None], module.mesh_spacing(), coord) + frame_id = np.zeros(coord.shape[0], dtype=np.int64) + energy = mesh_energy(charge, mesh, frame_id, module.kernel, module.pme_order) + return float(np.sum(energy)) + + +def test_pme_energy_converges_to_direct_reference() -> None: + """The PME energy approaches the direct k-space sum as the mesh refines.""" + coord = 3.0 + np.random.default_rng(7).uniform(0.0, 1.4, (6, 3)) + box = np.array([[9.0, 0.4, 0.0], [0.0, 8.5, 0.3], [0.0, 0.0, 9.5]]) + charge = np.array([[1.0], [-0.5], [1.0], [-0.5], [-0.5], [1.0]]) + charge = charge - charge.mean(axis=0, keepdims=True) + + kernel = LESKernel(alpha=0.3) + # converged direct k-space reference (tol 1e-10) + k_max = kernel.kspace_cutoff(1e-10) + rec = 2.0 * np.pi * np.linalg.inv(box) + n_d = np.ceil(k_max * np.linalg.norm(box, axis=1) / (2 * np.pi)).astype(int) + grids = np.meshgrid(*[np.arange(-n, n + 1) for n in n_d], indexing="ij") + n_all = np.stack(grids, -1).reshape(-1, 3) + k = n_all @ rec.T + k2 = (k * k).sum(-1) + keep = (k2 > 0) & (k2 <= k_max * k_max) + k, k2 = k[keep], k2[keep] + phase = k @ coord.T + rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge + s_sq = (np.abs(rho) ** 2).sum(-1) + k_hat = kernel.kspace(k2) + volume = abs(np.linalg.det(box)) + k0 = 2.0 * 0.3 / np.sqrt(np.pi) + e_ref = COULOMB_CONSTANT * ( + (k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum() + ) + errors = [] + for spacing in (2.0, None): + module = LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + kernel, + None, + periodic=True, + k_grid_mode="pme", + pme_spacing=spacing, + ) + errors.append(abs(_pme_energy(module, charge, coord, box) - e_ref)) + assert errors[1] < errors[0] # refinement improves the energy + # the default mesh reaches the accuracy it is chosen for + assert errors[1] < 1e-6 * abs(e_ref) + + +def test_pme_mesh_follows_the_cells_of_the_batch() -> None: + """The mesh of a batch resolves its cells at the spacing, without statistics.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + mesh_dims_of_cells, + ) + + long_range = _pme_long_range() + coord = np.zeros((9, 3)) + atype = np.zeros(9, dtype=np.int64) + small = _module_topologies(long_range, coord, np.array([4, 5]), _cells()[:1], atype) + large = _module_topologies( + long_range, coord, np.array([4, 5]), 2.0 * _cells(), atype + ) + spacing = long_range.mesh_spacing() + assert spacing == pytest.approx( + 0.75 * np.pi / long_range.kernel.kspace_cutoff(long_range.kspace_tolerance) + ) + for topologies, box in ((small, _cells()[:1]), (large, 2.0 * _cells())): + mesh = topologies["pme_mesh"] + dims = tuple(m.shape[0] for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)) + assert dims == mesh_dims_of_cells(box, spacing) + assert mesh.mask.all() + + +def test_pme_mode_configuration() -> None: + """The pme mode requires periodicity, and so do the mesh parameters.""" + with pytest.raises(ValueError, match="periodic"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + k_grid_mode="pme", + ) + for kwargs in ({"pme_spacing": 1.0}, {"pme_order": 6}): + with pytest.raises(ValueError, match="periodic"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + **kwargs, + ) + # every periodic mode takes them: a frozen model evaluates on the mesh + assert _periodic_long_range(pme_spacing=1.0).mesh_spacing() == 1.0 + assert _periodic_long_range().pme_order == 8 + with pytest.raises(ValueError, match="pme_order"): + _pme_long_range(pme_order=1) + with pytest.raises(ValueError, match="pme_spacing"): + _pme_long_range(pme_spacing=-1.0) + with pytest.raises(ValueError, match="k_grid_mode"): + _periodic_long_range(k_grid_mode="dense") + + +def test_pme_round_trip() -> None: + """The record keeps the pme mode, the order and the spacing.""" + long_range = _pme_long_range(pme_spacing=1.0, pme_order=5) + record = long_range.serialize() + assert record["k_grid_mode"] == "pme" + assert record["pme_order"] == 5 + assert record["pme_spacing"] == 1.0 + restored = LatentChargeLongRange.deserialize(record) + assert restored.k_grid_mode == "pme" + assert restored.pme_order == 5 + assert restored.pme_spacing == 1.0 + # the default order is recorded too, so that it survives a change of default + assert _periodic_long_range().serialize()["pme_order"] == 8 diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 9b61ee26cf..061ada901a 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -606,6 +606,9 @@ def has_default_chg_spin(self) -> bool: def get_default_chg_spin(self): return None + def has_long_range(self) -> bool: + return False + def _compose(self, children): return LinearEnergyAtomicModel(children, type_map=["Ni", "O"], weights="sum") diff --git a/source/tests/common/test_finetune_utils.py b/source/tests/common/test_finetune_utils.py index 92ad767271..622040c2d1 100644 --- a/source/tests/common/test_finetune_utils.py +++ b/source/tests/common/test_finetune_utils.py @@ -158,6 +158,28 @@ def test_finetune_rule_builder_updates_single_task_config_preserving_trainable() assert rule.get_has_new_type() +def test_finetune_rule_builder_carries_long_range_preserving_trainable(): + long_range = {"type": "latent_charge", "kernel": {"type": "les", "alpha": 0.7}} + pretrained = _model_config(["O", "H"]) + pretrained["long_range"] = deepcopy(long_range) + target = _model_config(["O", "H"], trainable=False) + target["long_range"] = { + "type": "latent_charge", + "charge_net": {"trainable": False}, + "kernel": {"type": "les", "alpha": 2.0, "trainable": False}, + } + + updated, _ = finetune.FinetuneRuleBuilder( + pretrained, + target, + change_model_params=True, + ).build() + + assert updated["long_range"]["kernel"]["alpha"] == 0.7 + assert updated["long_range"]["kernel"]["trainable"] is False + assert updated["long_range"]["charge_net"]["trainable"] is False + + def test_finetune_rule_builder_random_fitting_keeps_target_fitting_net(): pretrained = _model_config(["O", "H"], descriptor_sel=[8, 16], fitting_neuron=[32]) target = _model_config(["O", "H"], descriptor_sel=[1, 1], fitting_neuron=[2]) diff --git a/source/tests/infer/gen_dpa4c_long_range.py b/source/tests/infer/gen_dpa4c_long_range.py new file mode 100644 index 0000000000..8fa069eb05 --- /dev/null +++ b/source/tests/infer/gen_dpa4c_long_range.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate the DPA4C long-range .pt2 test models and reference values. + +A latent-charge long-range module sums over every pair of atoms of a complete +non-periodic system, and over the images of the cell of a periodic one when +it serves periodic systems, so the C++ interface hands the model what the sum +needs beyond the neighbor list: the graph lower of an uncompressed model +receives the pair list, or the mesh of the cell, that the interface builds +from the positions of the atoms, and the compact canonical lower of a +compressed model reads the positions and the cell themselves. One archive +covers each lower and each kind of module: + +- ``deeppot_dpa4c_long_range.pt2`` and + ``deeppot_dpa4c_long_range_periodic.pt2``: uncompressed, fp64, graph lower, + which compiles on the CPU; +- ``deeppot_dpa4c_long_range_compressed.pt2`` and + ``deeppot_dpa4c_long_range_periodic_compressed.pt2``: compressed, fp32, + compact canonical lower, whose operators are CUDA only; they are generated + for the CUDA variant of the C++ suite (``DP_VARIANT=cuda``). + +Every model conditions the descriptor on the charge state and constrains the +latent charges to the total charge that state names, so a state other than +the default changes the long-range energy as well as the descriptor. The +references are those of the Python interface on each archive. For a +non-periodic module they are taken without a cell, at the default state and +at ``[1, 2]`` (sections ``default`` and ``explicit``); for a periodic module, +in the cell :data:`BOX` and without a cell, at the neutral default state +(sections ``periodic`` and ``cluster``). The six atoms form two groups about +9 Å apart, beyond the cutoff of the descriptor, which only the long-range +module couples; in the cell, the first group meets its images across the +boundary. +""" + +import copy +import json +import os +import sys +import zipfile + +import numpy as np + +# The gen scripts run standalone, outside pytest's package machinery, so the +# relative imports the test suite uses do not apply here. +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from dpa4_fixtures import ( + jitter_zero_arrays, +) +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + +LONG_RANGE = { + "type": "latent_charge", + "channels": 1, + "charge_net": {"neuron": [8, 8], "seed": 9}, + "kernel": {"type": "les", "alpha": 0.8}, + "charge_constraint": {"source": "charge_spin"}, +} + +# The uncompressed model is fp64, so that its references hold to the +# precision of the double C++ interface. +CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 7, + "add_chg_spin_ebd": True, + "default_chg_spin": [0.0, 1.0], + }, + "fitting_net": {"neuron": [8, 8], "precision": "float64", "seed": 7}, + "long_range": LONG_RANGE, +} + +# Compression takes an fp32 descriptor and, for the compact canonical lower, +# a fitting that the fused fitting operator evaluates. +COMPRESSED_CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 8, + "precision": "float32", + "seed": 7, + "add_chg_spin_ebd": True, + "default_chg_spin": [0.0, 1.0], + }, + "fitting_net": { + "neuron": [32, 32], + "activation_function": "silu", + "resnet_dt": False, + "precision": "float32", + "seed": 7, + }, + "long_range": LONG_RANGE, +} + +# Jitters the zero-initialized arrays, among them the charge-state head of the +# descriptor, so that the charge state reaches the descriptor as well as the +# constraint. +_JITTER_SEED = 20261018 + +_ATYPES = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) +_COORDS = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ], + dtype=np.float64, +).reshape(1, 6, 3) +_OTHER_CHG_SPIN = np.array([[1.0, 2.0]], dtype=np.float64) +# A triclinic cell, row-major with the lattice vectors as rows. +BOX = np.array( + [[13.0, 0.0, 0.0], [0.6, 12.5, 0.0], [-0.4, 0.8, 13.5]], dtype=np.float64 +).reshape(1, 9) + + +def _periodic(config: dict) -> dict: + """The configuration with a long-range module that serves periodic systems.""" + config = copy.deepcopy(config) + config["long_range"]["periodic"] = True + return config + + +def _model_dict(config: dict) -> dict: + """Build the dpmodel of a configuration and jitter it in place of its zeros.""" + from deepmd.dpmodel.model.model import ( + get_model, + ) + + model = get_model(copy.deepcopy(config)) + return jitter_zero_arrays(model.serialize(), np.random.default_rng(_JITTER_SEED)) + + +def _check_metadata(pt2_path: str, lower: str, periodic: bool = False) -> None: + """Assert the archive declares its lower and the long-range input it needs.""" + with zipfile.ZipFile(pt2_path) as zf: + md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8")) + assert md["lower_input_kind"] == lower, md["lower_input_kind"] + expected = ["all_pairs", "pme_mesh"] if periodic else ["all_pairs"] + assert md["topologies"] == expected, md["topologies"] + if lower == "graph": + assert md["dim_chg_spin"] == 2 + else: + # The compressed descriptor folds the charge state; the archive + # carries the rebuild of every constant that holds it. + assert md["dim_chg_spin"] == 0 + assert len(md["charge_state_constants"]) == 5 + + +def _write_ref(pt2_path: str, ref_path: str, periodic: bool = False) -> None: + """Evaluate the reference cases on the archive and write the C++ sidecar.""" + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(pt2_path) + if periodic: + cases = (("periodic", BOX, None), ("cluster", None, None)) + else: + cases = (("default", None, None), ("explicit", None, _OTHER_CHG_SPIN)) + sections = {} + energies = {} + for name, box, charge_spin in cases: + condition = {} if charge_spin is None else {"charge_spin": charge_spin} + e, f, v, ae, av = dp.eval(_COORDS, box, _ATYPES, atomic=True, **condition) + assert np.all(np.isfinite(e)) and np.all(np.isfinite(f)) + energies[name] = float(e[0, 0]) + print(f"// {name} total energy: {e[0, 0]:.18e}") # noqa: T201 + sections[name] = { + "expected_e": ae[0, :, 0], + "expected_f": f[0], + "expected_v": av[0], + } + # The C++ cases compare the two cases, which needs a real difference. + first, second = (name for name, _, _ in cases) + assert abs(energies[second] - energies[first]) > 1e-3, ( + "the two cases must differ, or a call that ignored the cell or the " + "requested state would still match" + ) + write_expected_ref( + ref_path, + sections=sections, + source_script="source/tests/infer/gen_dpa4c_long_range.py", + ) + print(f"Wrote {ref_path}") # noqa: T201 + + +def _export_compressed(pt2_path: str, config: dict) -> None: + """Compress the fp32 model and export it for the CUDA device.""" + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + model = BaseModel.deserialize(_model_dict(config)) + model.atomic_model.enable_compression(0.5) + data = { + "model": model.serialize(), + "model_def_script": config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + print(f"Exporting to {pt2_path} ...") # noqa: T201 + deserialize_to_file(pt2_path, data, do_atomic_virial=True, lower_kind="auto") + + +def main() -> None: + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + ensure_inductor_compiler() + load_custom_ops() + + base_dir = os.path.dirname(__file__) + for suffix, periodic in (("", False), ("_periodic", True)): + config = _periodic(CONFIG) if periodic else CONFIG + stem = os.path.join(base_dir, f"deeppot_dpa4c_long_range{suffix}") + data = { + "model": _model_dict(config), + "model_def_script": config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + print(f"Exporting to {stem}.pt2 (lower_kind='graph') ...") # noqa: T201 + deserialize_to_file( + f"{stem}.pt2", data, do_atomic_virial=True, lower_kind="graph" + ) + _check_metadata(f"{stem}.pt2", "graph", periodic) + _write_ref(f"{stem}.pt2", f"{stem}.expected", periodic) + + if os.environ.get("DP_VARIANT") == "cuda": + compressed = f"{stem}_compressed" + compressed_config = ( + _periodic(COMPRESSED_CONFIG) if periodic else COMPRESSED_CONFIG + ) + _export_compressed(f"{compressed}.pt2", compressed_config) + _check_metadata(f"{compressed}.pt2", "dpa4c_canonical", periodic) + _write_ref(f"{compressed}.pt2", f"{compressed}.expected", periodic) + + print("\nDone!") # noqa: T201 + + +if __name__ == "__main__": + main() diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py index 4eef728456..eaa3b539a6 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -265,6 +265,22 @@ def test_serialization_preserves_parameters(self) -> None: original_parameters[name], ) + def test_state_dict_of_version_3_2_loads(self) -> None: + """A checkpoint holding the gather tables of DeePMD-kit 3.2 loads strictly.""" + state = self.descriptor.state_dict() + width = 32 + state["angular_channel_index"] = torch.arange(width) + state["angular_harmonic_index"] = torch.arange(width) + restored = self.build() + restored.load_state_dict(state) + torch.testing.assert_close( + self._evaluate(restored, self.coord), + self._evaluate(self.descriptor, self.coord), + ) + state["unknown"] = torch.zeros(1) + with pytest.raises(RuntimeError, match="unknown"): + restored.load_state_dict(state) + def test_deserialization_preserves_version_one_raw_radial_basis(self) -> None: data = self.descriptor.serialize() radial_basis = data["radial_basis"] diff --git a/source/tests/pt_expt/kernels/test_long_range_sum.py b/source/tests/pt_expt/kernels/test_long_range_sum.py new file mode 100644 index 0000000000..b61c7390c9 --- /dev/null +++ b/source/tests/pt_expt/kernels/test_long_range_sum.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Correctness of the operator of the latent-charge long range. + +For non-periodic frames the CUDA kernel evaluates every pair in float32 from +positions relative to the first node of each frame, carried as float32 head +and tail, and sums in compensated float32 before a float64 reduction; the +total-charge constraint and the per-frame reductions sum in float64 in a fixed +order. It is held to the float64 evaluation over the explicit pairs, at a +tolerance relative to the scale of each output, and to the exact invariants of +a pair sum: padding nodes and frames do not interact, the forces of a frame +cancel, and the energy, force, charge gradient and virial follow from one +another. + +For periodic frames the CUDA kernels evaluate the mesh sum in float32 on +B-spline weights formed in float64, and are held to the float64 mesh +evaluation of the training path, differentiated by autograd. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.long_range import ( + COULOMB_CONSTANT, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + mesh_dims_of_cells, +) +from deepmd.pt_expt.kernels.long_range import ( + _reference, + kernel_operator_arguments, + latent_charge_sum, + op_available, +) +from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + SOGKernel, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not op_available(), + reason="the long-range operator is CUDA only", +) + +KERNELS = { + "les": lambda: LESKernel(alpha=0.8), + "sog": lambda: SOGKernel( + amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2], precision="float32" + ), +} + + +def _system( + n_node: list[int], channels: int, padding: int, seed: int, offset: float +) -> tuple[torch.Tensor, ...]: + """Frames of random nodes far from the origin, with trailing padding nodes. + + Every seventh node carries no charge, and every frame a total charge. + """ + rng = np.random.default_rng(seed) + total = sum(n_node) + padding + coord = offset + rng.uniform(0.0, 4.0 * max(n_node) ** (1 / 3), (total, 3)) + charge = rng.normal(0.0, 0.5, (total, channels)) + device = torch.device("cuda") + node_mask = torch.ones(total, dtype=torch.bool, device=device) + node_mask[::7] = False + return ( + torch.tensor(coord, dtype=torch.float64, device=device), + torch.tensor(charge, dtype=torch.float32, device=device), + node_mask, + torch.tensor(n_node, dtype=torch.int64, device=device), + torch.tensor( + rng.integers(-2, 3, len(n_node)), dtype=torch.float64, device=device + ), + ) + + +def _reference_of( + system: tuple[torch.Tensor, ...], + kernel: Any, + channel: int, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, +) -> tuple[torch.Tensor, ...]: + """The float64 reference of the operator.""" + coord, charge, node_mask, n_node, total_charge = system + code, coefficient = kernel_operator_arguments(kernel) + if box is None: + box = torch.zeros(n_node.shape[0], 9, dtype=torch.float64) + return _reference( + coord, + charge.double(), + node_mask, + n_node, + total_charge, + box, + coefficient, + torch.zeros(1, dtype=torch.int64), + code, + channel, + COULOMB_CONSTANT, + order, + spacing, + ) + + +def _evaluate( + system: tuple[torch.Tensor, ...], + kernel: Any, + channel: int, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + actual = latent_charge_sum( + *system, + kernel, + channel, + COULOMB_CONSTANT, + box=box, + order=order, + spacing=spacing, + ) + return actual, _reference_of(system, kernel, channel, box, order, spacing) + + +NAMES = ("atom_energy", "charge_gradient", "force", "atom_virial", "virial") + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize( + ("n_node", "channels", "padding", "channel"), + [ + ([37, 1, 0, 300], 1, 5, 0), + ([37, 1, 0, 300], 1, 5, -1), + ([2048], 3, 0, -1), + ([5, 129], 8, 3, 2), + ], +) +def test_matches_float64_pair_sum( + kernel_name: str, n_node: list[int], channels: int, padding: int, channel: int +) -> None: + """Every output agrees with the float64 pair sum.""" + system = _system(n_node, channels, padding, 7, 1.0e3) + actual, expected = _evaluate(system, KERNELS[kernel_name](), channel) + _assert_close(actual, expected) + # Padding nodes join no frame. + if padding: + for value in actual[:4]: + assert torch.all(value[-padding:] == 0) + + +@pytest.mark.parametrize("alpha", [1e-3, 0.03, 0.3]) +@pytest.mark.parametrize("channel", [0, -1]) +def test_les_matches_float64_pair_sum_at_small_alpha_r( + alpha: float, channel: int +) -> None: + """The erf kernel keeps its precision where alpha r is small. + + The two terms of the slope of erf(x) / x agree to second order in x, so + that the slope of an approximate erf loses precision in proportion to + 1 / x^2 at atomic distances and a small alpha. + """ + system = _system([96], 2, 0, 5, 0.0) + _assert_close(*_evaluate(system, LESKernel(alpha=alpha), channel)) + + +def _assert_close( + actual: tuple[torch.Tensor, ...], + expected: tuple[torch.Tensor, ...], + tolerance: float = 2e-5, +) -> None: + """Every output agrees with the float64 reference, relative to its largest entry.""" + for name, value, reference in zip(NAMES, actual, expected, strict=True): + assert value.dtype == ( + torch.float64 if name == "atom_energy" else torch.float32 + ) + scale = reference.abs().max().item() + error = (value.double() - reference.double()).abs().max().item() + assert error <= tolerance * scale, (name, error, scale) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +def test_invariants(kernel_name: str) -> None: + """Frames are independent, forces cancel, and the virial is the pair virial.""" + n_node = [64, 200] + system = _system(n_node, 2, 0, 11, 0.0) + coord, charge, node_mask, n_node_t, total_charge = system + together = latent_charge_sum( + coord, + charge, + node_mask, + n_node_t, + total_charge, + KERNELS[kernel_name](), + 1, + COULOMB_CONSTANT, + ) + _, _, force, atom_virial, virial = together + # Each frame alone gives the same result. + first = n_node[0] + alone = latent_charge_sum( + coord[:first], + charge[:first], + node_mask[:first], + n_node_t[:1], + total_charge[:1], + KERNELS[kernel_name](), + 1, + COULOMB_CONSTANT, + ) + for both, single in zip(together, alone, strict=True): + torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0) + # Internal forces sum to zero per frame, and the virial of a frame is + # sum_i r_i (x) f_i. + for frame, (begin, end) in enumerate(((0, first), (first, sum(n_node)))): + frame_force = force[begin:end].double() + assert frame_force.sum(0).abs().max() <= 1e-5 * frame_force.abs().max() + position = coord[begin:end] - coord[begin] + expected = torch.einsum("ni,nj->ij", position, frame_force) + scale = expected.abs().max().item() + for value in ( + atom_virial[begin:end].double().sum(0).reshape(3, 3), + virial[frame].double().reshape(3, 3), + ): + torch.testing.assert_close(value, expected, rtol=0, atol=2e-5 * scale) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize("channel", [-1, 1]) +def test_reference_outputs_are_energy_derivatives( + kernel_name: str, channel: int +) -> None: + """The float64 force and charge gradient differentiate the energy.""" + coord, charge, node_mask, n_node, total_charge = _system([9, 6], 2, 2, 3, 0.0) + coord = coord.clone().requires_grad_(True) + charge = charge.double().requires_grad_(True) + atom_energy, charge_gradient, force, _, _ = _reference_of( + (coord, charge, node_mask, n_node, total_charge), + KERNELS[kernel_name](), + channel, + ) + grad_coord, grad_charge = torch.autograd.grad(atom_energy.sum(), (coord, charge)) + torch.testing.assert_close(force, -grad_coord, rtol=1e-9, atol=1e-9) + torch.testing.assert_close( + charge_gradient.detach(), grad_charge, rtol=1e-9, atol=1e-9 + ) + + +def test_bitwise_reproducible() -> None: + """Two evaluations of the same input agree bit for bit, with or without cells.""" + system = _system([3000], 1, 0, 5, 0.0) + for box in (None, _cells(1, 40.0)): + first = latent_charge_sum( + *system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box + ) + second = latent_charge_sum( + *system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box + ) + for a, b in zip(first, second, strict=True): + assert torch.equal(a, b) + + +# --------------------------------------------------------------------------- +# Periodic frames +# --------------------------------------------------------------------------- + + +def _cells(frame_count: int, length: float, seed: int = 3) -> torch.Tensor: + """Triclinic cells of about ``length`` per side, row-major, one row per frame.""" + rng = np.random.default_rng(seed) + cells = [] + for _ in range(frame_count): + cell = np.diag(length * rng.uniform(0.8, 1.2, 3)) + cell[1, 0], cell[2, 0], cell[2, 1] = length * rng.uniform(-0.2, 0.2, 3) + cells.append(cell.reshape(9)) + return torch.tensor(np.asarray(cells), dtype=torch.float64) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize( + ("n_node", "channels", "padding", "channel", "order"), + [ + ([37, 1, 300], 1, 5, 0, 8), + ([37, 1, 300], 1, 5, -1, 4), + ([500], 3, 0, -1, 6), + ([5, 129], 8, 3, 2, 8), + ], +) +def test_mesh_matches_float64_reference( + kernel_name: str, + n_node: list[int], + channels: int, + padding: int, + channel: int, + order: int, +) -> None: + """Every output of a periodic frame agrees with the mesh evaluation of training. + + The positions extend beyond the cell, so that the fractional coordinates + wrap. + """ + system = _system(n_node, channels, padding, 7, -5.0) + box = _cells(len(n_node), 12.0) + actual, expected = _evaluate( + system, KERNELS[kernel_name](), channel, box, order, 0.6 + ) + _assert_close(actual, expected, 1e-6) + if padding: + for value in actual[:4]: + assert torch.all(value[-padding:] == 0) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +def test_mesh_virial_is_the_strain_derivative(kernel_name: str) -> None: + """The virial of a periodic frame is the strain derivative of its energy. + + Straining the positions and the cell together by ``I + eps`` changes the + energy by ``-W : eps``; the atomic virial of the charged nodes sums to W. + The central difference takes the float64 mesh energy, whose rounding stays + far below that of the float32 outputs of the operator. + """ + system = _system([60], 2, 0, 13, 0.0) + coord, charge, node_mask, n_node, total_charge = system + box = _cells(1, 10.0) + kernel = KERNELS[kernel_name]() + _, _, _, atom_virial, virial = latent_charge_sum( + *system, kernel, 0, COULOMB_CONSTANT, box=box, spacing=0.5 + ) + h = 1e-5 + finite = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + energies = [] + for sign in (1.0, -1.0): + strain = torch.eye(3, dtype=torch.float64) + strain[a, b] += sign * h + strained = ( + coord @ strain.T.to(coord.device), + charge, + node_mask, + n_node, + total_charge, + ) + cell = box.reshape(3, 3) @ strain.T + # Both strained cells keep the mesh of the cell. + assert mesh_dims_of_cells(cell.reshape(1, 9).numpy(), 0.5) == ( + mesh_dims_of_cells(box.numpy(), 0.5) + ) + energies.append( + _reference_of(strained, kernel, 0, cell.reshape(1, 9), 8, 0.5)[ + 0 + ].sum() + ) + finite[a, b] = -(energies[0] - energies[1]).item() / (2 * h) + scale = np.abs(finite).max() + np.testing.assert_allclose( + virial[0].double().cpu().numpy().reshape(3, 3), finite, atol=2e-6 * scale + ) + np.testing.assert_allclose( + atom_virial.double().sum(0).cpu().numpy().reshape(3, 3), + finite, + atol=2e-6 * scale, + ) + + +@pytest.mark.parametrize("order", [2, 3, 5, 7, 9, 10, 11, 12]) +def test_mesh_orders(order: int) -> None: + """Every B-spline order agrees with the mesh evaluation of training. + + The atomic energies of these small cells are a fraction of the self + energy, to a few 1e-7 of which the float32 mesh resolves them. + """ + system = _system([40, 23], 2, 3, 5, -5.0) + box = _cells(2, 9.0) + actual, expected = _evaluate(system, KERNELS["sog"](), 1, box, order, 0.7) + _assert_close(actual, expected, 1e-5) + + +def test_mesh_stencil_wider_than_the_axis() -> None: + """A stencil wider than an axis of the mesh wraps around it more than once.""" + system = _system([12], 1, 0, 9, 0.0) + box = _cells(1, 3.0) + # Four points per axis, against eight of the stencil. + assert mesh_dims_of_cells(box.numpy(), 1.0) == (4, 4, 4) + actual, expected = _evaluate(system, KERNELS["les"](), 0, box, 8, 1.0) + _assert_close(actual, expected, 1e-6) + + +@pytest.mark.parametrize("periodic", [False, True]) +def test_frame_without_charged_nodes(periodic: bool) -> None: + """A frame whose nodes carry no charge contributes zeros, the others are unaffected.""" + n_node = [40, 6, 50] + coord, charge, node_mask, n_node_t, total_charge = _system(n_node, 2, 0, 5, 0.0) + node_mask[40:46] = False + system = (coord, charge, node_mask, n_node_t, total_charge) + box = _cells(3, 10.0) if periodic else None + actual, expected = _evaluate(system, KERNELS["les"](), 0, box, 8, 0.6) + _assert_close(actual, expected, 1e-6) + for value in actual[:4]: + assert torch.all(value[40:46] == 0) + assert torch.all(actual[4][1] == 0) + + +def test_mesh_frames_are_independent() -> None: + """Every periodic frame is evaluated on its own mesh, bit for bit.""" + n_node = [64, 200] + system = _system(n_node, 2, 0, 11, 0.0) + coord, charge, node_mask, n_node_t, total_charge = system + box = _cells(2, 9.0) + kernel = KERNELS["sog"]() + together = latent_charge_sum( + *system, kernel, 1, COULOMB_CONSTANT, box=box, spacing=0.5 + ) + first = n_node[0] + alone = latent_charge_sum( + coord[:first], + charge[:first], + node_mask[:first], + n_node_t[:1], + total_charge[:1], + kernel, + 1, + COULOMB_CONSTANT, + box=box[:1], + spacing=0.5, + ) + for both, single in zip(together, alone, strict=True): + torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0) + + +def test_mixed_frames_are_rejected() -> None: + """A call holds periodic frames or non-periodic ones, not both.""" + system = _system([10, 10], 1, 0, 1, 0.0) + box = _cells(2, 9.0) + box[1] = 0.0 + with pytest.raises(RuntimeError, match="periodic"): + latent_charge_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py new file mode 100644 index 0000000000..c48dc0ae4c --- /dev/null +++ b/source/tests/pt_expt/model/test_long_range.py @@ -0,0 +1,2002 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Long-range latent-charge modules on the DPA4C graph lower (pt_expt).""" + +import copy + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.long_range import ( + COULOMB_CONSTANT, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + build_pair_graph, +) +from deepmd.infer import ( + DeepPot, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.train.training import ( + _CompiledModel, + _get_model_structure_key, + model_change_out_bias, +) +from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, +) +from deepmd.pt_expt.utils import ( + env, +) +from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, +) + +from ..compile_utils import ( + REQUIRES_SUPPORTED_COMPILE, +) + +KERNELS = { + "les": {"type": "les", "alpha": 0.8}, + "sog": {"type": "sog", "amplitude": [0.7, -0.2, 0.1], "width": [0.8, 1.7, 3.2]}, +} +CHANNELS = {"les": 1, "sog": 3} +NEUTRAL = {"source": "neutral"} + + +def _config(kernel: str = "les", constraint: dict | None = NEUTRAL) -> dict: + """Small float64 DPA4C model with a latent-charge long-range module.""" + return { + "type_map": ["A", "B"], + "descriptor": { + "type": "dpa4c", + "rcut": 3.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 17, + }, + "fitting_net": { + "type": "ener", + "neuron": [16, 16], + "precision": "float64", + "seed": 19, + }, + "long_range": { + "type": "latent_charge", + "channels": CHANNELS[kernel], + "charge_net": {"neuron": [8, 8], "precision": "float64", "seed": 23}, + "kernel": copy.deepcopy(KERNELS[kernel]), + "charge_constraint": copy.deepcopy(constraint), + }, + } + + +def _model(kernel: str = "les", **kwargs: object) -> torch.nn.Module: + return get_model(_config(kernel, **kwargs)).to(env.DEVICE).eval() + + +def _tensor(x: np.ndarray, dtype: torch.dtype = torch.float64) -> torch.Tensor: + return torch.tensor(x, dtype=dtype, device=env.DEVICE) + + +def _batch(seed: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + """Three frames of 6, 3 and 5 real atoms, padded to a common width.""" + rng = np.random.default_rng(seed) + coord = rng.uniform(0.0, 6.0, (3, 6, 3)) + atype = np.array( + [[0, 1, 0, 1, 1, 0], [1, 0, 0, -1, -1, -1], [0, 0, 1, 1, 0, -1]], + dtype=np.int64, + ) + return _tensor(coord), _tensor(atype, torch.int64) + + +def _set_constant_charges(model: torch.nn.Module, charges: list[float]) -> None: + """Zero every network weight and give the atoms of each type a fixed charge. + + The energy is then the long-range pair sum alone, of charges that do not + depend on the geometry; the kernel keeps its configured parameters. + """ + with torch.no_grad(): + for name, parameter in model.atomic_model.named_parameters(): + if ".kernel." not in f".{name}": + parameter.zero_() + bias = model.atomic_model.long_range.charge_net.bias_atom_e + bias.zero_() + bias[:, 0] = torch.as_tensor(charges, dtype=bias.dtype) + + +def _reference_energy(coord: np.ndarray, charge: np.ndarray, kernel: str) -> float: + """Direct pair sum of the latent-charge energy of one frame, in eV.""" + energy = 0.0 + for i in range(len(coord)): + for j in range(i + 1, len(coord)): + r = np.linalg.norm(coord[i] - coord[j]) + if kernel == "les": + from scipy.special import ( + erf, + ) + + k = erf(0.8 * r) / r + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k = np.sum(amplitude * np.exp(-(r**2) / (2 * width**2))) + energy += COULOMB_CONSTANT * float(charge[i] @ charge[j]) * k + return energy + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_energy_matches_direct_pair_sum(kernel: str) -> None: + """With fixed charges the energy is the documented pair sum, attributed half per atom.""" + model = _model(kernel, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + coord = np.array([[[0.0, 0.0, 0.0], [1.3, 0.0, 0.0], [0.0, 2.1, 0.4]]]) + atype = np.array([[0, 1, 1]]) + out = model(_tensor(coord), _tensor(atype, torch.int64)) + charge = np.zeros((3, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, -0.5] + expected = _reference_energy(coord[0], charge, kernel) + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-12) + np.testing.assert_allclose( + out["atom_energy"].sum().item(), expected, rtol=1e-12, atol=1e-14 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_force_virial_hessian_match_finite_differences(kernel: str) -> None: + """Force, affine-strain virial and Hessian are derivatives of the complete energy.""" + model = _model(kernel) + # Atoms within the cutoff of each other make the charges depend on the + # geometry, so the derivatives hold the response of the charges. + rng = np.random.default_rng(3) + coord = rng.uniform(0.0, 2.5, (1, 4, 3)) + atype = np.array([[0, 1, 0, 1]]) + coord_t, atype_t = _tensor(coord), _tensor(atype, torch.int64) + out = model(coord_t, atype_t) + + def energy(c: np.ndarray) -> float: + return model(_tensor(c), atype_t)["energy"].item() + + h = 1e-5 + force_fd = np.zeros((4, 3)) + for i in range(4): + for d in range(3): + plus, minus = coord.copy(), coord.copy() + plus[0, i, d] += h + minus[0, i, d] -= h + force_fd[i, d] = -(energy(plus) - energy(minus)) / (2 * h) + np.testing.assert_allclose( + out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-8 + ) + + # The virial is -dE/d(strain) at zero strain, for x -> (1 + eps) x. + virial_fd = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + strain = np.zeros((3, 3)) + strain[a, b] = h + plus = coord @ (np.eye(3) + strain) + minus = coord @ (np.eye(3) - strain) + virial_fd[a, b] = -(energy(plus) - energy(minus)) / (2 * h) + np.testing.assert_allclose( + out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-8 + ) + + hessian_model = _model(kernel) + hessian_model.enable_hessian() + hessian = hessian_model(coord_t, atype_t)["hessian"].detach().cpu().numpy() + hessian = hessian.reshape(12, 12) + hessian_fd = np.zeros((12, 12)) + for k in range(12): + plus, minus = coord.copy(), coord.copy() + plus.reshape(-1)[k] += h + minus.reshape(-1)[k] -= h + f_plus = model(_tensor(plus), atype_t)["force"].detach().cpu().numpy() + f_minus = model(_tensor(minus), atype_t)["force"].detach().cpu().numpy() + hessian_fd[:, k] = -(f_plus - f_minus).reshape(-1) / (2 * h) + np.testing.assert_allclose(hessian, hessian_fd, atol=1e-6) + np.testing.assert_allclose(hessian, hessian.T, atol=1e-10) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_atomic_reductions_and_padding(kernel: str) -> None: + """Atomic energy and virial sum to the frame values; padding atoms are inert.""" + model = _model(kernel) + coord, atype = _batch() + out = model(coord, atype, do_atomic_virial=True) + torch.testing.assert_close(out["atom_energy"].sum(1), out["energy"]) + torch.testing.assert_close(out["atom_virial"].sum(1), out["virial"]) + assert out["force"][1, 3:].abs().max() == 0 + moved = coord.clone() + moved[1, 3:] += 50.0 + moved_out = model(moved, atype) + torch.testing.assert_close(moved_out["energy"], out["energy"], rtol=0, atol=0) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_frames_are_independent(kernel: str) -> None: + """A frame keeps its energy and force whatever other frames are batched with it.""" + model = _model(kernel) + coord, atype = _batch() + batched = model(coord, atype) + for frame, nloc in enumerate((6, 3, 5)): + alone = model(coord[frame : frame + 1, :nloc], atype[frame : frame + 1, :nloc]) + torch.testing.assert_close(alone["energy"][0], batched["energy"][frame]) + torch.testing.assert_close( + alone["force"][0], batched["force"][frame, :nloc], atol=1e-12, rtol=0 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_ragged_matches_rectangular(kernel: str) -> None: + """The flat ragged entry evaluates the same model as the padded one.""" + model = _model(kernel) + coord, atype = _batch() + padded = model(coord, atype) + keep = atype.reshape(-1) >= 0 + ragged = model.forward_ragged( + coord.reshape(-1, 3)[keep], + atype.reshape(-1)[keep], + torch.tensor([6, 3, 5], device=env.DEVICE), + ) + torch.testing.assert_close(ragged["energy"], padded["energy"]) + torch.testing.assert_close(ragged["force"], padded["force"].reshape(-1, 3)[keep]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_empty_and_single_atom_frames(kernel: str) -> None: + """Frames without pairs contribute no long-range energy and stay finite.""" + model = _model(kernel) + coord = _tensor(np.random.default_rng(5).uniform(0.0, 3.0, (3, 2, 3))) + atype = _tensor(np.array([[0, -1], [1, 0], [-1, -1]]), torch.int64) + out = model(coord, atype) + assert torch.isfinite(out["energy"]).all() + assert torch.isfinite(out["force"]).all() + single = model(coord[:1, :1], atype[:1, :1]) + torch.testing.assert_close(single["energy"][0], out["energy"][0]) + + +def test_permutation_and_rigid_motion_invariance() -> None: + """The energy is invariant under atom permutation, translation and rotation.""" + model = _model("sog") + coord, atype = _batch() + coord, atype = coord[:1], atype[:1] + out = model(coord, atype) + perm = torch.tensor([3, 0, 5, 1, 4, 2], device=env.DEVICE) + permuted = model(coord[:, perm], atype[:, perm]) + torch.testing.assert_close(permuted["energy"], out["energy"]) + torch.testing.assert_close(permuted["force"], out["force"][:, perm]) + rotation, _ = torch.linalg.qr( + _tensor(np.random.default_rng(13).normal(size=(3, 3))) + ) + moved = model(coord @ rotation.T + 3.7, atype) + torch.testing.assert_close(moved["energy"], out["energy"]) + torch.testing.assert_close(moved["force"], out["force"] @ rotation.T) + + +@pytest.mark.parametrize("source", ["charge_spin", "fparam"]) +def test_constraint_default_matches_explicit_target(source: str) -> None: + """An omitted frame condition takes its default for the constraint as well.""" + config = _config("les", constraint={"source": source, "fparam_index": 0}) + if source == "charge_spin": + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [2.0, 1.0] + explicit = {"charge_spin": _tensor(np.array([[2.0, 1.0], [2.0, 1.0]]))} + other = {"charge_spin": _tensor(np.array([[0.0, 1.0], [0.0, 1.0]]))} + else: + config["fitting_net"]["numb_fparam"] = 1 + config["fitting_net"]["default_fparam"] = [2.0] + explicit = {"fparam": _tensor(np.array([[2.0], [2.0]]))} + other = {"fparam": _tensor(np.array([[0.0], [0.0]]))} + model = get_model(config).to(env.DEVICE).eval() + coord, atype = _batch() + coord, atype = coord[:2], atype[:2] + implicit_out = model(coord, atype) + explicit_out = model(coord, atype, **explicit) + torch.testing.assert_close(implicit_out["energy"], explicit_out["energy"]) + torch.testing.assert_close(implicit_out["force"], explicit_out["force"]) + other_out = model(coord, atype, **other) + assert (other_out["energy"] - implicit_out["energy"]).abs().max() > 1e-6 + + +@pytest.mark.parametrize("constraint", [None, NEUTRAL]) +def test_pair_exclusion_removes_the_interaction(constraint: dict | None) -> None: + """Pairs of excluded types do not interact; their atoms keep their charges.""" + config = _config("les", constraint=constraint) + config["pair_exclude_types"] = [[0, 1]] + model = get_model(config).to(env.DEVICE).eval() + _set_constant_charges(model, [1.0, -1.0]) + coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]])) + atype = _tensor(np.array([[0, 0, 1]]), torch.int64) + out = model(coord, atype) + # Only the two type-0 atoms interact. The neutral constraint shifts the + # charges (1, 1, -1) of the three atoms to (2/3, 2/3, -4/3). + charge = 1.0 if constraint is None else 2.0 / 3.0 + erf = torch.erf(torch.tensor(3.2, dtype=torch.float64)).item() + expected = COULOMB_CONSTANT * charge**2 * erf / 4.0 + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-12) + assert out["force"][0, 2].abs().max() == 0 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_les_kernel_derivatives_down_to_contact(dtype: torch.dtype) -> None: + """The radial derivatives of LES keep the precision of the dtype near contact. + + The derivatives of erf(alpha r) / r subtract terms that agree to second + order in alpha r, which the series of the kernel avoids; the reference + differentiates in 40-digit arithmetic. + """ + import mpmath + + from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + ) + + alpha = 0.8 + kernel = LESKernel(alpha=alpha) + r = torch.tensor( + np.geomspace(1e-5, 6.0, 80), dtype=dtype, device=env.DEVICE + ).requires_grad_(True) + (first,) = torch.autograd.grad(kernel(r * r).sum(), r, create_graph=True) + (second,) = torch.autograd.grad(first.sum(), r) + + def les(t: mpmath.mpf) -> mpmath.mpf: + return mpmath.erf(alpha * t) / t + + rtol = 2e-5 if dtype == torch.float32 else 1e-11 + for order, actual in ((1, first), (2, second)): + with mpmath.workdps(40): + expected = np.array( + [ + float(mpmath.diff(les, mpmath.mpf(float(t)), order)) + for t in r.detach().cpu() + ] + ) + # The second derivative changes sign near 1.2 Å, where it is the + # difference of terms much larger than itself. + np.testing.assert_allclose( + actual.detach().cpu().numpy(), + expected, + rtol=rtol, + atol=rtol * np.abs(expected).max(), + err_msg=str(order), + ) + + +def test_owned_frames_across_domains() -> None: + """Halo nodes are rejected on a single domain and kept across domains.""" + from deepmd.pt_expt.model.make_model import ( + check_owned_frames, + ) + + n_node = torch.tensor([5, 3], dtype=torch.int64) + n_local = torch.tensor([4, 3], dtype=torch.int64) + check_owned_frames(n_node, n_node.clone()) + check_owned_frames(n_node, n_local, torch.ones(1, dtype=torch.int64)) + for communicator in (None, torch.zeros(1, dtype=torch.int64)): + with pytest.raises(ValueError, match="n_local == n_node"): + check_owned_frames(n_node, n_local, communicator) + + +def test_graph_lower_rejects_halo_nodes() -> None: + """A graph whose frames carry halo nodes is rejected, not summed without them. + + The halo nodes of a graph without communication would carry no charge and + remove their interactions from the sum of a frame held whole. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_topologies, + ) + from deepmd.pt_expt.utils.graph_builder import ( + build_neighbor_graph_for_method, + ) + + model = _model("les", constraint=None) + coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]])) + atype = _tensor(np.array([[0, 0, 1]]), torch.int64) + graph = build_neighbor_graph_for_method( + "dense", coord, atype, None, model.get_rcut(), None + ) + topologies = build_topologies( + model.atomic_model.required_topologies(), + coord.reshape(-1, 3), + graph.n_node, + atype=atype.reshape(-1), + ) + inputs = ( + atype.reshape(-1), + graph.n_node, + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + ) + + def lower(n_local: torch.Tensor) -> dict[str, torch.Tensor]: + return model.forward_common_lower_graph( + inputs[0], inputs[1], n_local, *inputs[2:], topologies=topologies + ) + + lower(graph.n_node.clone()) + with pytest.raises(ValueError, match="n_local == n_node"): + lower(graph.n_node - 1) + + +def test_linear_composition_takes_the_exclusions_of_the_composition() -> None: + """A long-range child follows the pair exclusion of its composition. + + The children of a composition share its neighbor graph and pair topology, + so the exclusions of the composition give the child the semantics of the + same model alone, and an exclusion of the child itself is rejected. + """ + child = _config("les", constraint=None) + child["type"] = "standard" + composition = { + "type": "linear_ener", + "type_map": child["type_map"], + "weights": "sum", + "models": [child], + "pair_exclude_types": [[0, 1]], + } + model = get_model(copy.deepcopy(composition)).to(env.DEVICE).eval() + with torch.no_grad(): + for name, parameter in model.atomic_model.named_parameters(): + if ".kernel." not in f".{name}": + parameter.zero_() + bias = model.atomic_model.models[0].long_range.charge_net.bias_atom_e + bias[:, 0] = torch.as_tensor([1.0, -1.0], dtype=bias.dtype) + coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]])) + atype = _tensor(np.array([[0, 0, 1]]), torch.int64) + # As in test_pair_exclusion_removes_the_interaction: only the two type-0 + # atoms interact. + erf = torch.erf(torch.tensor(3.2, dtype=torch.float64)).item() + np.testing.assert_allclose( + model(coord, atype)["energy"].item(), + COULOMB_CONSTANT * erf / 4.0, + rtol=1e-12, + ) + for key, value, error, match in ( + ("pair_exclude_types", [[0, 1]], ValueError, "from the composition"), + ("atom_exclude_types", [1], NotImplementedError, "atom_exclude_types"), + ): + config = copy.deepcopy(composition) + config.pop("pair_exclude_types") + config["models"][0][key] = value + with pytest.raises(error, match=match): + get_model(config) + # a periodic module shares its composition with no other long-range module + periodic = copy.deepcopy(child) + periodic["long_range"]["periodic"] = True + for models in ([periodic, child], [periodic, periodic]): + config = copy.deepcopy(composition) + config.pop("pair_exclude_types") + config["models"] = copy.deepcopy(models) + with pytest.raises(NotImplementedError, match="serves periodic systems"): + get_model(config) + config["models"] = [copy.deepcopy(child), copy.deepcopy(child)] + get_model(config) + # a periodic module takes no pair exclusion, alone or in a composition + config = copy.deepcopy(composition) + config["models"] = [copy.deepcopy(periodic)] + with pytest.raises(NotImplementedError, match="pair_exclude_types"): + get_model(config) + config = copy.deepcopy(periodic) + config["type"] = "standard" + config["pair_exclude_types"] = [[0, 1]] + with pytest.raises(NotImplementedError, match="pair_exclude_types"): + get_model(config) + + +@pytest.mark.parametrize("compiled", [False, True]) +def test_vacuum_reference_with_long_range(compiled: bool) -> None: + """An isolated atom keeps its reference energy; reference rows join no pair.""" + config = _config("les", constraint={"source": "charge_spin"}) + config["type_map"] = ["O", "H"] + config["preset_out_bias"] = {"energy": {"O": -3.0, "H": 0.5}} + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [0.0, 1.0] + config["fitting_net"]["vacuum_ref"] = True + model = get_model(config).to(env.DEVICE).eval() + assert model.get_fitting_net().needs_vacuum_descriptor() + rng = np.random.default_rng(11) + coord = _tensor(rng.uniform(0.0, 2.5, (2, 4, 3))) + atype = _tensor(np.array([[0, 1, 1, 0], [1, 0, -1, -1]]), torch.int64) + if compiled: + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model(coord, atype) + out = runner(coord, atype) + torch.testing.assert_close(out["energy"], reference["energy"]) + torch.testing.assert_close(out["force"], reference["force"]) + return + out = model(coord, atype) + for frame in range(2): + alone = model(coord[frame : frame + 1], atype[frame : frame + 1]) + torch.testing.assert_close(alone["energy"][0], out["energy"][frame]) + # The output statistics set the bias; an atom without neighbors and in + # its reference state then has exactly the bias of its type. + bias = model.atomic_model.out_bias.clone() + bias[0, :, 0] = torch.tensor([-3.0, 0.5]) + model.atomic_model.set_out_bias(bias) + isolated = model( + coord[:1, :1], + atype[:1, :1], + charge_spin=_tensor(np.array([[0.0, 3.0]])), + ) + torch.testing.assert_close( + isolated["energy"].reshape(()), + torch.tensor(-3.0, dtype=torch.float64, device=env.DEVICE), + rtol=0, + atol=1e-12, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_self_label_bias_adjustment_is_invariant(kernel: str) -> None: + """Labelling a sample with the model's own energies leaves its bias unchanged.""" + model = _model(kernel) + with torch.no_grad(): + model.atomic_model.long_range.charge_net.bias_atom_e[:, 0] = torch.tensor( + [0.8, -0.6] + ) + coord, atype = _batch() + coord, atype = coord[:1], atype[:1] + energy = model(coord, atype)["energy"].detach().cpu().numpy() + natoms = np.array([[6, 6, 3, 3]], dtype=np.int32) + sample = [ + { + "coord": coord.cpu().numpy(), + "atype": atype.cpu().numpy(), + "box": None, + "natoms": natoms, + "energy": energy.reshape(1, 1), + "find_energy": np.float32(1.0), + } + ] + bias_before = model.atomic_model.out_bias.detach().clone() + model_change_out_bias(model, lambda: sample, "change-by-statistic") + torch.testing.assert_close( + model.atomic_model.out_bias, bias_before, atol=1e-10, rtol=0 + ) + after = model(coord, atype)["energy"].detach().cpu().numpy() + np.testing.assert_allclose(after, energy, atol=1e-10) + + +@pytest.mark.parametrize("weights", ["sum", "mean", [0.3, 1.4]]) +def test_linear_composition_combines_complete_energies(weights: object) -> None: + """A linear composition weights the complete energy of each child. + + The second child is an ordinary model, which reads no topology. + """ + child = _config("les") + child["type"] = "standard" + other = copy.deepcopy(child) + other.pop("long_range") + other["descriptor"]["seed"] = 31 + composition = ( + get_model( + { + "type": "linear_ener", + "type_map": child["type_map"], + "weights": weights, + "models": [child, other], + } + ) + .to(env.DEVICE) + .eval() + ) + standalone = [_model("les"), get_model(other).to(env.DEVICE).eval()] + for index, model in enumerate(standalone): + model.atomic_model.load_state_dict( + composition.atomic_model.models[index].state_dict() + ) + coord, atype = _batch() + out = composition(coord, atype) + if weights == "sum": + w = [1.0, 1.0] + elif weights == "mean": + w = [0.5, 0.5] + else: + w = weights + energy = sum( + ww * m(coord, atype)["energy"] for ww, m in zip(w, standalone, strict=True) + ) + force = sum( + ww * m(coord, atype)["force"] for ww, m in zip(w, standalone, strict=True) + ) + torch.testing.assert_close(out["energy"], energy) + torch.testing.assert_close(out["force"], force) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_parameter_gradients_and_freezing(kernel: str) -> None: + """Energy and force losses reach the charge network and the kernel; frozen parts stay out. + + Every parameter of the module moves the losses, except under a constraint + the common charge bias of the constrained channel, whose gradient + vanishes (see test_constraint_removes_the_common_charge_bias). + """ + coord, atype = _batch() + for constraint in (None, NEUTRAL): + model = _model(kernel, constraint=constraint).train() + out = model(coord, atype) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + long_range = model.atomic_model.long_range + bias = long_range.charge_net.nets[0].layers[-1].b + for name, parameter in long_range.named_parameters(): + assert parameter.grad is not None, name + assert torch.isfinite(parameter.grad).all(), name + gradient = parameter.grad + if constraint is not None and parameter is bias: + # Channel 0 is the constrained channel. + gradient = gradient[..., 1:] + if gradient.numel(): + assert gradient.abs().max() > 0, name + + config = _config(kernel) + config["long_range"]["kernel"]["trainable"] = False + config["long_range"]["charge_net"]["trainable"] = False + frozen = get_model(config).to(env.DEVICE).train() + names = [ + name + for name, parameter in frozen.named_parameters() + if parameter.requires_grad and "long_range" in name + ] + assert names == [] + frozen(coord, atype)["energy"].sum().backward() + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_constraint_removes_the_common_charge_bias(kernel: str) -> None: + """A uniform shift of the raw charges of the constrained channel changes nothing. + + The constraint subtracts the mean shift of every frame, so the output bias + of that channel leaves energy and force unchanged, while that of another + channel does not. + """ + model = _model(kernel) + coord, atype = _batch() + reference = model(coord, atype) + bias = model.atomic_model.long_range.charge_net.nets[0].layers[-1].b + for channel in range(bias.shape[-1]): + with torch.no_grad(): + bias[..., channel] += 0.3 + out = model(coord, atype) + with torch.no_grad(): + bias[..., channel] -= 0.3 + if channel == 0: + torch.testing.assert_close(out["energy"], reference["energy"]) + torch.testing.assert_close(out["force"], reference["force"]) + else: + assert (out["energy"] - reference["energy"]).abs().max() > 1e-6 + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_serialization_round_trip(kernel: str) -> None: + """The portable record and the state dict restore the same model.""" + model = _model(kernel) + rng = np.random.default_rng(5) + with torch.no_grad(): + for parameter in model.atomic_model.long_range.parameters(): + parameter.add_(0.01 * _tensor(rng.normal(size=tuple(parameter.shape)))) + coord, atype = _batch() + expected = model(coord, atype) + record = model.serialize() + assert record["@version"] == 3 + restored = type(model).deserialize(copy.deepcopy(record)).to(env.DEVICE).eval() + torch.testing.assert_close(restored(coord, atype)["energy"], expected["energy"]) + rebuilt = _model(kernel) + rebuilt.load_state_dict(model.state_dict()) + torch.testing.assert_close(rebuilt(coord, atype)["force"], expected["force"]) + + +def test_construction_is_reproducible() -> None: + """Without a seed of its own the charge network follows the fitting seed.""" + config = _config("les") + del config["long_range"]["charge_net"]["seed"] + first, second = (get_model(copy.deepcopy(config)) for _ in range(2)) + for (name, a), (_, b) in zip( + first.named_parameters(), second.named_parameters(), strict=True + ): + torch.testing.assert_close(a, b, rtol=0, atol=0, msg=name) + + +def test_ordinary_record_keeps_version() -> None: + """A model without a long-range module keeps its record and constructor.""" + config = _config("les") + config.pop("long_range") + model = get_model(config) + record = model.serialize() + assert record["@version"] == 2 + assert "long_range" not in record + assert not model.atomic_model.has_long_range() + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_checkpoint_inference_matches_direct_forward(tmp_path, kernel: str) -> None: + """DeepPot on a checkpoint builds the topology over the real atoms only.""" + model = _model(kernel) + checkpoint = tmp_path / f"long_range_{kernel}.pt" + wrapper = ModelWrapper(model, model_params=_config(kernel)) + torch.save({"model": wrapper.state_dict()}, checkpoint) + coord, atype = _batch(seed=7) + expected = model(coord, atype, do_atomic_virial=True) + dp = DeepPot(str(checkpoint), auto_batch_size=False) + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + coord_np = coord.cpu().numpy() + actual = dict( + zip( + names, + dp.eval(coord_np, None, atype.cpu().numpy(), atomic=True, mixed_type=True), + strict=True, + ) + ) + for name in names: + np.testing.assert_allclose( + actual[name], + expected[name].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + coord_np[1, 3:] += 40.0 + moved = dp.eval(coord_np, None, atype.cpu().numpy(), mixed_type=True) + np.testing.assert_allclose(moved[0], actual["energy"], rtol=0, atol=1e-12) + + +@pytest.mark.parametrize("grid_mode", ["ragged", "pme"]) +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_checkpoint_inference_matches_direct_forward( + tmp_path, kernel: str, grid_mode: str +) -> None: + """DeepPot on a periodic checkpoint builds the periodic topology and matches.""" + if grid_mode == "pme": + model = _periodic_pme_model(kernel) + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "pme" + else: + model = _periodic_model(kernel) + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "ragged" + checkpoint = tmp_path / f"periodic_{kernel}_{grid_mode}.pt" + wrapper = ModelWrapper(model, model_params=config) + torch.save({"model": wrapper.state_dict()}, checkpoint) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + expected = model(coord_t, atype_t, box=box_t, do_atomic_virial=True) + dp = DeepPot(str(checkpoint), auto_batch_size=False) + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + actual = dict( + zip( + names, + dp.eval( + coord, + box, + atype, + atomic=True, + mixed_type=True, + ), + strict=True, + ) + ) + for name in names: + np.testing.assert_allclose( + actual[name], + expected[name].detach().cpu().numpy(), + rtol=1e-9, + atol=1e-9, + err_msg=name, + ) + + +@pytest.mark.parametrize( + ("kernel", "source", "suffix"), + [("les", "neutral", ".pt2"), ("sog", "neutral", ".pte"), ("les", "fparam", ".pte")], +) +def test_frozen_model_matches_direct_forward( + tmp_path, kernel: str, source: str, suffix: str +) -> None: + """A frozen model reads the pair topology as named inputs of any length.""" + config = _config(kernel, constraint={"source": source, "fparam_index": 0}) + frame = {} + if source == "fparam": + config["fitting_net"]["numb_fparam"] = 1 + frame["fparam"] = np.array([[1.0], [-0.5], [2.0]]) + model = get_model(config).to(env.DEVICE).eval() + frozen = tmp_path / f"long_range_{kernel}{suffix}" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="graph", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + assert dp.deep_eval.metadata["topologies"] == ["all_pairs"] + if suffix == ".pte": + # The pair topology follows the thirteen inputs of the graph lower. + signature = torch.export.load(str(frozen)).graph_signature + assert len(signature.user_inputs) == 13 + 3 + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + for seed, nloc in ((7, 6), (8, 4)): + coord, atype = _batch(seed=seed) + coord, atype = coord[:, :nloc], atype[:, :nloc] + expected = model( + coord, + atype, + do_atomic_virial=True, + **{key: _tensor(value) for key, value in frame.items()}, + ) + actual = dp.eval( + coord.cpu().numpy(), + None, + atype.cpu().numpy(), + atomic=True, + mixed_type=True, + **frame, + ) + for name, value in zip(names, actual, strict=True): + np.testing.assert_allclose( + value, + expected[name].detach().cpu().numpy().reshape(value.shape), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + + +@pytest.mark.parametrize( + ("kernel", "source"), [("les", None), ("les", "neutral"), ("sog", "charge_spin")] +) +def test_coordinate_evaluation_matches_pair_graph( + kernel: str, source: str | None +) -> None: + """The evaluation from node positions matches the pair-graph evaluation. + + It returns the atomic energy, the derivative of the energy with respect to + the descriptor through the constrained charges, and the force at a fixed + descriptor, all of which the pair graph yields by automatic + differentiation. Two frames, an uncharged node and trailing padding nodes + exercise the frame sums of the constraint. + """ + from deepmd.pt_expt.kernels.long_range import ( + op_available, + ) + + if not op_available(): + pytest.skip("the long-range operator is not available") + config = _config(kernel, constraint=None if source is None else {"source": source}) + if source == "charge_spin": + config["descriptor"]["add_chg_spin_ebd"] = True + long_range = get_model(config).atomic_model.long_range.eval() + if source == "charge_spin": + long_range.fold_charge_state([2.0, 1.0]) + rng = np.random.default_rng(7) + n_node = torch.tensor([9, 6], device=env.DEVICE) + node_mask = torch.ones(17, dtype=torch.bool, device=env.DEVICE) + node_mask[4] = False + node_mask[15:] = False + coord = _tensor(rng.uniform(0.0, 6.0, (17, 3))) + descriptor = _tensor(rng.normal(size=(17, long_range.charge_net.dim_descrpt))) + atype = _tensor(rng.integers(0, 2, 17), torch.int64) + + atom_energy, descriptor_gradient, force, _, _ = long_range.energy_force_from_coord( + descriptor, atype, node_mask, n_node, coord + ) + + leaf_descriptor = descriptor.clone().requires_grad_(True) + leaf_coord = coord.clone().requires_grad_(True) + pairs = build_pair_graph(leaf_coord, n_node, atype, node_mask) + expected = long_range.call_graph( + leaf_descriptor, atype, node_mask, n_node, {ALL_PAIRS: pairs} + ) + expected_descriptor_gradient, expected_coord_gradient = torch.autograd.grad( + expected.sum(), (leaf_descriptor, leaf_coord) + ) + for actual, target in ( + (atom_energy, expected), + (descriptor_gradient, expected_descriptor_gradient), + (force, -expected_coord_gradient), + ): + target = target.detach() + scale = target.abs().max().item() + assert scale > 0.0 + np.testing.assert_allclose( + actual.detach().to(torch.float64).cpu().numpy(), + target.cpu().numpy(), + rtol=0.0, + atol=2e-5 * scale, + ) + + +@pytest.mark.parametrize( + ("kernel", "source"), [("les", "charge_spin"), ("sog", "neutral")] +) +def test_compressed_frozen_model_matches_uncompressed( + tmp_path, kernel: str, source: str +) -> None: + """A compressed model keeps its long-range module through the freeze. + + On a CUDA device the freeze selects the compact canonical lower, whose + long-range operator reads the positions of the atoms; on the CPU the + graph lower reads the pair lists. A constraint that reads the total + charge from the charge state is folded with the descriptor and follows + the state the deployment serves. Both match the uncompressed model. + """ + config = _config(kernel, constraint={"source": source}) + config["descriptor"].update(n_radial=8, precision="float32") + if source == "charge_spin": + config["descriptor"].update(add_chg_spin_ebd=True, default_chg_spin=[0.0, 1.0]) + config["fitting_net"].update( + neuron=[32, 32], activation_function="silu", resnet_dt=False + ) + config["fitting_net"]["precision"] = "float32" + config["long_range"]["charge_net"] = {"neuron": [16, 16], "seed": 23} + model = get_model(config).to(env.DEVICE).eval() + generator = torch.Generator().manual_seed(5) + with torch.no_grad(): + for parameter in model.parameters(): + noise = torch.randn(parameter.shape, generator=generator) + parameter.add_(0.05 * noise.to(parameter)) + reference = copy.deepcopy(model) + model.atomic_model.enable_compression(0.5) + frozen = tmp_path / f"long_range_compressed_{kernel}.pt2" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="auto", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + expected_lower = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph" + assert dp.deep_eval.metadata["lower_input_kind"] == expected_lower + rng = np.random.default_rng(3) + coord = rng.uniform(0.0, 7.0, (1, 40, 3)) + 100.0 + atype = rng.integers(0, 2, (1, 40)) + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + # The tables of the compressed descriptor approximate its networks. On the + # graph lower the charges of this model amplify that error at a charged + # state, to about 1e-4 of the scale at [1, 2]; the compact canonical lower + # stays within a few 1e-6. + tolerance = 2e-5 if expected_lower == "dpa4c_canonical" else 5e-4 + states = [[0.0, 1.0], [1.0, 2.0], [0.0, 1.0]] if source == "charge_spin" else [None] + for state in states: + condition = {} if state is None else {"charge_spin": np.array([state])} + expected = reference( + _tensor(coord), + _tensor(atype, torch.int64), + do_atomic_virial=True, + **{key: _tensor(value) for key, value in condition.items()}, + ) + actual = dp.eval(coord, None, atype, atomic=True, **condition) + for name, value in zip(names, actual, strict=True): + target = expected[name].detach().cpu().numpy().reshape(value.shape) + scale = np.abs(target).max() + assert np.abs(value - target).max() <= tolerance * scale + 1e-5, ( + name, + state, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the complete model, train and eval. + + In training the gradient of an energy and force loss, whose force part is + a second-order derivative of the energy, reaches the long-range + parameters as in eager execution. + """ + model = _model(kernel) + coord, atype = _batch() + parameters = list(model.atomic_model.long_range.parameters()) + + def loss(out: dict) -> torch.Tensor: + return out["energy"].square().sum() + out["force"].square().sum() + + for training in (True, False): + model.train(training) + compiled = _CompiledModel( + model, _get_model_structure_key(model), compile_eval=True + ) + compiled.train(training) + reference = model(coord, atype) + result = compiled(coord, atype) + assert compiled._compiled_lower_by_mode, "the compiled graph was not used" + torch.testing.assert_close(result["energy"], reference["energy"]) + torch.testing.assert_close(result["force"], reference["force"]) + torch.testing.assert_close(result["virial"], reference["virial"]) + if training: + expected = torch.autograd.grad(loss(reference), parameters) + actual = torch.autograd.grad(loss(result), parameters) + for expected_grad, actual_grad in zip(expected, actual, strict=True): + torch.testing.assert_close(actual_grad, expected_grad) + + +def test_unsupported_combinations_fail() -> None: + """Unsupported inputs, routes, combinations and descriptors are rejected.""" + model = _model("les") + coord, atype = _batch() + box = 10.0 * torch.eye(3, dtype=torch.float64, device=env.DEVICE).reshape(1, 9) + with pytest.raises(NotImplementedError, match="non-periodic"): + model(coord[:1], atype[:1], box=box) + nlist = torch.full((1, 6, 4), -1, dtype=torch.int64, device=env.DEVICE) + with pytest.raises(NotImplementedError, match="NeighborGraph"): + model.atomic_model.forward_atomic(coord[:1], atype[:1], nlist) + config = _config("les") + config["spin"] = {"use_spin": [True, False], "scheme": "native"} + with pytest.raises(NotImplementedError): + get_model(config) + config = _config("les") + config["use_srtab"] = "table.txt" + config["sw_rmin"], config["sw_rmax"] = 0.5, 1.0 + with pytest.raises(NotImplementedError, match="ZBL"): + get_model(config) + config = _config("les") + config["type_map"] = ["O", "H"] + config["bridging_method"] = "ZBL" + with pytest.raises(NotImplementedError, match="bridging"): + get_model(config) + config = _config("les") + config["atom_exclude_types"] = [1] + with pytest.raises(NotImplementedError, match="atom_exclude_types"): + get_model(config) + child = _config("les") + child["type"] = "standard" + with pytest.raises(NotImplementedError, match="atom_exclude_types"): + get_model( + { + "type": "linear_ener", + "type_map": child["type_map"], + "weights": "sum", + "models": [child], + "atom_exclude_types": [1], + } + ) + config = _config("les") + config["descriptor"] = { + "type": "se_e2_a", + "rcut": 3.0, + "rcut_smth": 2.0, + "sel": [8, 8], + } + with pytest.raises(NotImplementedError, match="DPA4C"): + get_model(config) + + +def test_compressing_the_descriptor_alone_is_rejected() -> None: + """A folded descriptor and a constraint reading the condition are rejected.""" + config = _config("les", constraint={"source": "charge_spin"}) + config["descriptor"].update( + n_radial=8, + precision="float32", + add_chg_spin_ebd=True, + default_chg_spin=[0.0, 1.0], + ) + config["fitting_net"]["precision"] = "float32" + model = get_model(config).to(env.DEVICE).eval() + model.atomic_model.descriptor.enable_compression(0.5) + coord, atype = _batch() + with pytest.raises(ValueError, match="folds the charge/spin condition"): + model(coord[:1], atype[:1]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_dpmodel_reference_matches(kernel: str) -> None: + """The dpmodel reference evaluates the same energy from the portable record.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + model = _model(kernel) + coord, atype = _batch() + expected = model(coord, atype) + reference = BaseModelDP.deserialize(model.serialize()) + ret = reference.call_common( + coord.cpu().numpy(), atype.cpu().numpy(), neighbor_graph_method="dense" + ) + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-12, + atol=1e-12, + ) + np.testing.assert_allclose( + ret["energy"].reshape(3, 6), + expected["atom_energy"].detach().cpu().numpy().reshape(3, 6), + rtol=1e-12, + atol=1e-12, + ) + + +def test_pair_graph_skips_padding() -> None: + """The pair topology joins real atoms of one frame only and ends in masked guards.""" + coord = torch.arange(15, dtype=torch.float64).reshape(5, 3) + pairs = build_pair_graph( + coord, + torch.tensor([3, 2]), + torch.tensor([0, -1, 1, 0, 1]), + torch.tensor([True, False, True, True, True]), + ) + real = pairs.index[:, pairs.mask] + assert real.tolist() == [[0, 3], [2, 4]] + assert (~pairs.mask).sum() == 2 + + +# --------------------------------------------------------------------------- +# Periodic (k-space) evaluation +# --------------------------------------------------------------------------- + + +def _periodic_model(kernel: str, **kwargs: object) -> torch.nn.Module: + """A periodic model over the explicit wave numbers of every frame.""" + config = _config(kernel, **kwargs) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "ragged" + return get_model(config).to(env.DEVICE).eval() + + +def _periodic_case() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """A compact cluster (all atoms within rcut of each other) in a large triclinic cell.""" + box = np.array([[9.0, 0.4, 0.0], [0.0, 8.5, 0.3], [0.0, 0.0, 9.5]]) + rng = np.random.default_rng(7) + coord = 3.0 + rng.uniform(0.0, 1.4, (6, 3)) + atype = np.array([0, 1, 0, 1, 1, 0], dtype=np.int64) + return coord[None], atype[None], box[None] + + +def _kspace_reference( + coord: np.ndarray, charge: np.ndarray, box: np.ndarray, kernel: str +) -> float: + """Independent k-space evaluation of the periodic long-range energy, in eV. + + E = k_e [ 1/(2V) sum_{k != 0} K_hat(k) |rho(k)|^2 - K(0)/2 sum_i |q_i|^2 ]. + """ + rec = 2.0 * np.pi * np.linalg.inv(box) # columns g_d: k(n) = rec @ n + # the module's default kspace_tolerance=1e-8 sets the truncation + if kernel == "les": + k_max = 2.0 * 0.8 * np.sqrt(np.log(1e8)) + else: + k_max = np.sqrt(2.0 * np.log(1e8)) / 0.8 + n_d = np.ceil(k_max * np.linalg.norm(box, axis=1) / (2 * np.pi)).astype(int) + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=float) for n in n_d], indexing="ij" + ) + n_all = np.stack(grids, axis=-1).reshape(-1, 3) + k = n_all @ rec.T + k2 = np.sum(k * k, axis=-1) + keep = (k2 > 0) & (k2 <= k_max * k_max) + k, k2 = k[keep], k2[keep] + phase = k @ coord.T # (nk, N) + rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge # (nk, C) + s_sq = np.abs(rho) ** 2 + s_sq = s_sq.sum(axis=-1) + if kernel == "les": + k_hat = 4.0 * np.pi * np.exp(-k2 / (4.0 * 0.8**2)) / k2 + k0 = 2.0 * 0.8 / np.sqrt(np.pi) + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k_hat = np.sum( + amplitude[None, :] + * (2.0 * np.pi) ** 1.5 + * width[None, :] ** 3 + * np.exp(-0.5 * k2[:, None] * width[None, :] ** 2), + axis=-1, + ) + k0 = amplitude.sum() + volume = abs(np.linalg.det(box)) + return float( + COULOMB_CONSTANT + * ((k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum()) + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_energy_matches_kspace_reference(kernel: str) -> None: + """With fixed charges the periodic energy is the documented k-space sum.""" + model = _periodic_model(kernel, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + coord, atype, box = _periodic_case() + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + charge = np.zeros((6, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] + expected = _kspace_reference(coord[0], charge, box[0], kernel) + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-9, atol=1e-9) + np.testing.assert_allclose( + out["atom_energy"].sum().item(), expected, rtol=1e-9, atol=1e-9 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_force_virial_hessian_match_finite_differences(kernel: str) -> None: + """Force, strain virial and Hessian derive from the complete periodic energy.""" + model = _periodic_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + out = model(coord_t, atype_t, box=box_t) + + def energy(c: np.ndarray, b: np.ndarray) -> float: + return model(_tensor(c), atype_t, box=_tensor(b))["energy"].item() + + h = 1e-5 + force_fd = np.zeros((6, 3)) + for i in range(6): + for d in range(3): + plus, minus = coord.copy(), coord.copy() + plus[0, i, d] += h + minus[0, i, d] -= h + force_fd[i, d] = -(energy(plus, box) - energy(minus, box)) / (2 * h) + np.testing.assert_allclose( + out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-6 + ) + + # The virial is -dE/d(strain) at zero strain, for r -> (I + eps) r and + # B -> B (I + eps)^T (the strain family of the k-space strain leaf). + virial_fd = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + strain = np.zeros((3, 3)) + strain[a, b] = h + plus = coord @ (np.eye(3) + strain).T + minus = coord @ (np.eye(3) - strain).T + box_plus = box @ (np.eye(3) + strain).T + box_minus = box @ (np.eye(3) - strain).T + virial_fd[a, b] = -(energy(plus, box_plus) - energy(minus, box_minus)) / ( + 2 * h + ) + np.testing.assert_allclose( + out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-5 + ) + + hessian_model = _periodic_model(kernel) + hessian_model.load_state_dict(model.state_dict()) + hessian_model.enable_hessian() + hessian = ( + hessian_model(coord_t, atype_t, box=box_t)["hessian"].detach().cpu().numpy() + ) + hessian = hessian.reshape(18, 18) + hessian_fd = np.zeros((18, 18)) + for kk in range(18): + plus, minus = coord.copy(), coord.copy() + plus.reshape(-1)[kk] += h + minus.reshape(-1)[kk] -= h + f_plus = ( + model(_tensor(plus), atype_t, box=box_t)["force"].detach().cpu().numpy() + ) + f_minus = ( + model(_tensor(minus), atype_t, box=box_t)["force"].detach().cpu().numpy() + ) + hessian_fd[:, kk] = -(f_plus - f_minus).reshape(-1) / (2 * h) + np.testing.assert_allclose(hessian, hessian_fd, atol=1e-4) + np.testing.assert_allclose(hessian, hessian.T, atol=1e-10) + + +@pytest.mark.parametrize("grid_mode", ["ragged", "pme"]) +def test_trace_samples_give_every_axis_its_own_length(grid_mode: str) -> None: + """No two dynamic axes of the trace-time topologies share a length, nor take 2 or 3.""" + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_topologies, + ) + + model = ( + _periodic_model("les") if grid_mode == "ragged" else _periodic_pme_model("les") + ) + nframes, n_node, n_pair = 5, 7, 11 + topologies = build_synthetic_topologies( + model, n_node, n_pair, dtype=torch.float64, device=env.DEVICE, nframes=nframes + ) + lengths = {("", "frame"): nframes, ("", "node"): n_node} + for name, topology in topologies.items(): + for field, axes in type(topology).AXES.items(): + for axis, label in axes.items(): + key = ("", label) if label in ("frame", "node") else (name, label) + size = getattr(topology, field).shape[axis] + assert lengths.setdefault(key, size) == size, (name, field, label) + sizes = list(lengths.values()) + assert len(set(sizes)) == len(sizes), lengths + assert not {2, 3} & set(sizes), lengths + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_periodic_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the periodic model like the eager one.""" + model = _periodic_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model.eval()(coord_t, atype_t, box=box_t) + model.train() + out = runner(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8) + # the backward reaches the charge network and the kernel through the + # compiled graph + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] + assert grads and all(grad is not None for grad in grads) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_model_serves_non_periodic_batches(kernel: str) -> None: + """A periodic-enabled model keeps its non-periodic evaluation.""" + model = _periodic_model(kernel) + plain = _model(kernel) + plain.load_state_dict(model.state_dict()) + coord, atype = _batch() + periodic_enabled = model(coord, atype) + plain_out = plain(coord, atype) + torch.testing.assert_close(periodic_enabled["energy"], plain_out["energy"]) + torch.testing.assert_close(periodic_enabled["force"], plain_out["force"]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_frames_are_independent(kernel: str) -> None: + """Each frame's k grid and energy are independent of the rest of the batch.""" + model = _periodic_model(kernel) + coord, atype, box = _periodic_case() + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + atype2 = np.concatenate([atype, atype], axis=0) + box2 = np.concatenate([box, box * 1.1], axis=0) + batched = model(_tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2)) + for frame in range(2): + alone = model( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_serialization_round_trip(kernel: str) -> None: + """The record keeps the periodic flag and evaluates identically.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + model = _periodic_model(kernel) + assert model.atomic_model.required_topologies() == ("all_pairs", "k_grid") + record = model.serialize() + assert record["long_range"]["periodic"] is True + reference = BaseModelDP.deserialize(record) + assert reference.atomic_model.required_topologies() == ("all_pairs", "k_grid") + coord, atype, box = _periodic_case() + expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + ret = reference.call_common( + coord, + atype, + box, + neighbor_graph_method="dense", + ) + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-9, + atol=1e-9, + ) + + +def _periodic_frames() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Two frames of the periodic case, the second strained and shifted.""" + coord, atype, box = _periodic_case() + strain = np.array([[1.04, 0.02, 0.0], [0.0, 0.97, 0.01], [0.03, 0.0, 1.02]]) + coord2 = np.concatenate([coord, coord @ strain.T + 0.3], axis=0) + box2 = np.concatenate([box, box @ strain.T], axis=0) + return coord2, np.concatenate([atype, atype], axis=0), box2 + + +@pytest.mark.parametrize( + ("kernel", "grid_mode", "suffix"), + [("les", "ragged", ".pt2"), ("sog", "pme", ".pte"), ("sog", "fixed", ".pt2")], +) +def test_periodic_frozen_model_matches_mesh_evaluation( + tmp_path, kernel: str, grid_mode: str, suffix: str +) -> None: + """A frozen periodic model evaluates every periodic frame on the mesh of its cell. + + Whatever its training mode, the frozen model reads the mesh of the cells + of a batch as inputs of dynamic size, and matches the model evaluated on + the mesh; a batch without a cell keeps the direct sum. + """ + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = grid_mode + if grid_mode == "fixed": + # Bounds that hold every wave number of the cells within the cutoff. + config["long_range"]["k_grid_bounds"] = [14, 14, 14] + model = get_model(config).to(env.DEVICE).eval() + frozen = tmp_path / f"periodic_{kernel}_{grid_mode}{suffix}" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="graph", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + assert dp.deep_eval.metadata["topologies"] == ["all_pairs", "pme_mesh"] + mesh_model = copy.deepcopy(model) + mesh_model.atomic_model.long_range.use_mesh_evaluation() + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + coord, atype, box = _periodic_frames() + for frame in (slice(0, 1), slice(0, 2)): + expected = mesh_model( + _tensor(coord[frame]), + _tensor(atype[frame], torch.int64), + box=_tensor(box[frame]), + do_atomic_virial=True, + ) + actual = dp.eval( + coord[frame], box[frame], atype[frame], atomic=True, mixed_type=True + ) + for name, value in zip(names, actual, strict=True): + np.testing.assert_allclose( + value, + expected[name].detach().cpu().numpy().reshape(value.shape), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + # The mesh approximates the sum of the training mode. + trained = model( + _tensor(coord[:1]), _tensor(atype[:1], torch.int64), box=_tensor(box[:1]) + ) + np.testing.assert_allclose( + dp.eval(coord[:1], box[:1], atype[:1], mixed_type=True)[0], + trained["energy"].detach().cpu().numpy(), + rtol=1e-5, + ) + cluster, cluster_type = _batch() + expected = model(cluster, cluster_type) + actual = dp.eval( + cluster.cpu().numpy(), None, cluster_type.cpu().numpy(), mixed_type=True + ) + np.testing.assert_allclose( + actual[0], expected["energy"].detach().cpu().numpy(), rtol=1e-10, atol=1e-10 + ) + + +@pytest.mark.parametrize(("kernel", "source"), [("les", "neutral"), ("sog", None)]) +def test_periodic_compressed_frozen_model_matches_uncompressed( + tmp_path, kernel: str, source: str | None +) -> None: + """A compressed periodic model keeps the mesh evaluation through the freeze. + + On a CUDA device the compact canonical lower evaluates every periodic + frame on the mesh of its cell with the long-range operator; on the CPU + the graph lower reads the mesh. Both match the uncompressed model on the + mesh, for a periodic frame and for a cluster. + """ + config = _config(kernel, constraint=None if source is None else {"source": source}) + config["descriptor"].update(n_radial=8, precision="float32") + config["fitting_net"].update( + neuron=[32, 32], activation_function="silu", resnet_dt=False + ) + config["fitting_net"]["precision"] = "float32" + config["long_range"]["charge_net"] = {"neuron": [16, 16], "seed": 23} + config["long_range"]["periodic"] = True + model = get_model(config).to(env.DEVICE).eval() + generator = torch.Generator().manual_seed(5) + with torch.no_grad(): + for parameter in model.parameters(): + noise = torch.randn(parameter.shape, generator=generator) + parameter.add_(0.05 * noise.to(parameter)) + reference = copy.deepcopy(model) + reference.atomic_model.long_range.use_mesh_evaluation() + model.atomic_model.enable_compression(0.5) + frozen = tmp_path / f"periodic_compressed_{kernel}.pt2" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="auto", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + expected_lower = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph" + assert dp.deep_eval.metadata["lower_input_kind"] == expected_lower + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + rng = np.random.default_rng(3) + box = np.array([[[12.0, 0.0, 0.0], [0.8, 11.0, 0.0], [-0.5, 0.6, 13.0]]]) + periodic = rng.uniform(0.0, 1.0, (1, 40, 3)) @ box[0] + cluster = rng.uniform(0.0, 7.0, (1, 40, 3)) + 100.0 + atype = rng.integers(0, 2, (1, 40)) + for coord, cell in ((periodic, box), (cluster, None)): + expected = reference( + _tensor(coord), + _tensor(atype, torch.int64), + box=None if cell is None else _tensor(cell), + do_atomic_virial=True, + ) + actual = dp.eval(coord, cell, atype, atomic=True) + for name, value in zip(names, actual, strict=True): + target = expected[name].detach().cpu().numpy().reshape(value.shape) + scale = np.abs(target).max() + # The tables of the compressed descriptor approximate its + # networks, to a few 1e-6 of the scale of every output. + assert np.abs(value - target).max() <= 2e-5 * scale + 1e-5, ( + name, + cell is None, + np.abs(value - target).max(), + scale, + ) + # A cell without volume marks no system, on either lower: the graph lower + # meets it in the neighbor search, the compact lower at its input. + with pytest.raises((ValueError, RuntimeError), match=r"no volume|singular"): + dp.eval(periodic, np.zeros((1, 9)), atype) + + +# --------------------------------------------------------------------------- +# Periodic fixed k-grid mode +# --------------------------------------------------------------------------- + + +def _periodic_fixed_model( + kernel: str, boxes: np.ndarray | None = None, **kwargs: object +) -> torch.nn.Module: + """A periodic model in the fixed k-grid mode, sized by the given cells.""" + config = _config(kernel, **kwargs) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "fixed" + model = get_model(config).to(env.DEVICE).eval() + if boxes is not None: + boxes = torch.as_tensor(np.asarray(boxes, dtype=np.float64).reshape(-1, 9)) + model.atomic_model.long_range.compute_input_stats([{"box": boxes}]) + return model + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_matches_ragged(kernel: str) -> None: + """The fixed grid with its mask evaluates like the ragged enumeration.""" + coord, atype, box = _periodic_case() + ragged = _periodic_model(kernel) + fixed = _periodic_fixed_model(kernel, box) + # the state dicts of the two grid modes differ by the fixed grid buffer + fixed.load_state_dict(ragged.state_dict(), strict=False) + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + expected = ragged(coord_t, atype_t, box=box_t) + got = fixed(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(got[key], expected[key], rtol=1e-9, atol=1e-12) + # the gradient reaches the charge network and the kernel without NaN, + # though masked points carry a singular kernel transform nowhere + fixed.train() + out = fixed(coord_t, atype_t, box=box_t) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [ + parameter.grad + for parameter in fixed.atomic_model.long_range.parameters() + if parameter.grad is not None + ] + assert grads and all(torch.isfinite(grad).all() for grad in grads) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_mixed_cells(kernel: str) -> None: + """One shared grid serves a batch of different cells, each masked to its sphere.""" + coord, atype, box = _periodic_case() + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + atype2 = np.concatenate([atype, atype], axis=0) + box2 = np.concatenate([box, box * 1.4], axis=0) + fixed = _periodic_fixed_model(kernel, box2) + ragged = _periodic_model(kernel) + ragged.load_state_dict(fixed.state_dict(), strict=False) + coord_t, atype_t, box_t = ( + _tensor(coord2), + _tensor(atype2, torch.int64), + _tensor(box2), + ) + batched = fixed(coord_t, atype_t, box=box_t) + expected = ragged(coord_t, atype_t, box=box_t) + torch.testing.assert_close( + batched["energy"], expected["energy"], rtol=1e-9, atol=1e-12 + ) + torch.testing.assert_close( + batched["force"], expected["force"], rtol=1e-9, atol=1e-12 + ) + # the batched energy of a frame equals its standalone evaluation + for frame in range(2): + alone = fixed( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10 + ) + + +def _kspace_reference_bounded( + coord: np.ndarray, + charge: np.ndarray, + box: np.ndarray, + kernel: str, + bounds: np.ndarray, +) -> float: + """The k-space reference of :func:`_kspace_reference` restricted to the + integer indices within ``bounds`` — the truncated sum of a fixed grid. + """ + rec = 2.0 * np.pi * np.linalg.inv(box) + if kernel == "les": + k_max = 2.0 * 0.8 * np.sqrt(np.log(1e8)) + else: + k_max = np.sqrt(2.0 * np.log(1e8)) / 0.8 + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=float) for n in bounds], indexing="ij" + ) + n_all = np.stack(grids, axis=-1).reshape(-1, 3) + k = n_all @ rec.T + k2 = np.sum(k * k, axis=-1) + keep = (k2 > 0) & (k2 <= k_max * k_max) + k, k2 = k[keep], k2[keep] + phase = k @ coord.T + rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge + s_sq = (np.abs(rho) ** 2).sum(axis=-1) + if kernel == "les": + k_hat = 4.0 * np.pi * np.exp(-k2 / (4.0 * 0.8**2)) / k2 + k0 = 2.0 * 0.8 / np.sqrt(np.pi) + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k_hat = np.sum( + amplitude[None, :] + * (2.0 * np.pi) ** 1.5 + * width[None, :] ** 3 + * np.exp(-0.5 * k2[:, None] * width[None, :] ** 2), + axis=-1, + ) + k0 = amplitude.sum() + volume = abs(np.linalg.det(box)) + return float( + COULOMB_CONSTANT + * ((k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum()) + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_truncates_oversized_cells(kernel: str) -> None: + """A cell beyond the grid is evaluated with the grid alone, not an error.""" + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + big_box = box * 3.0 + with pytest.warns(UserWarning, match="truncated"): + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(big_box)) + bounds = model.atomic_model.long_range.k_grid_bounds.detach().cpu().numpy() + charge = np.zeros((6, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] + expected = _kspace_reference_bounded(coord[0], charge, big_box[0], kernel, bounds) + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_needs_statistics_or_bounds(kernel: str) -> None: + """Without statistics or explicit bounds the grid is not sized; bounds skip them.""" + coord, atype, box = _periodic_case() + unsized = _periodic_fixed_model(kernel) + with pytest.raises(ValueError, match="not sized yet"): + unsized(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "fixed" + config["long_range"]["k_grid_bounds"] = [12, 12, 12] + model = get_model(config).to(env.DEVICE).eval() + np.testing.assert_array_equal( + model.atomic_model.long_range.k_grid_bounds.cpu().numpy(), [12, 12, 12] + ) + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + reference = _periodic_model(kernel) + reference.load_state_dict(model.state_dict(), strict=False) + expected = reference(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + torch.testing.assert_close(out["energy"], expected["energy"], rtol=1e-9, atol=1e-12) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_restores_from_state_dict(kernel: str) -> None: + """The state dict of a sized grid loads into an unsized model, as a checkpoint does.""" + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box) + restored = _periodic_fixed_model(kernel) + restored.load_state_dict(model.state_dict()) + args = (_tensor(coord), _tensor(atype, torch.int64)) + expected = model(*args, box=_tensor(box)) + got = restored(*args, box=_tensor(box)) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(got[key], expected[key], rtol=1e-12, atol=1e-12) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_serialization_round_trip(kernel: str) -> None: + """The record keeps the sized grid and evaluates identically.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box) + record = model.serialize() + assert record["long_range"]["k_grid_mode"] == "fixed" + assert "k_grid_bounds" in record["long_range"] + reference = BaseModelDP.deserialize(record) + assert reference.atomic_model.long_range.k_grid_mode == "fixed" + expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + ret = reference.call_common(coord, atype, box, neighbor_graph_method="dense") + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-9, + atol=1e-9, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_periodic_fixed_grid_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the fixed-grid model like the eager one.""" + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box) + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model.eval()(coord_t, atype_t, box=box_t) + model.train() + out = runner(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] + assert grads and all(grad is not None for grad in grads) + + +# --------------------------------------------------------------------------- +# Periodic decoupled k grid +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_decoupled_grid_ignores_kernel_drift(kernel: str) -> None: + """A decoupled model selects the k points with the initial kernel but + evaluates them with the live one: the drift moves the energy, not the grid. + """ + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "ragged" + config["long_range"]["decouple_k_grid"] = True + model = get_model(config).to(env.DEVICE).eval() + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + long_range = model.atomic_model.long_range + grid_args = ( + coord_t.reshape(-1, 3), + torch.tensor([6], dtype=torch.int64, device=env.DEVICE), + box_t, + ) + build = model.atomic_model.build_topologies + before = build(*grid_args, atype=atype_t.reshape(-1))["k_grid"] + energy_before = model(coord_t, atype_t, box=box_t)["energy"] + # move the kernel far enough that a coupled grid would change + with torch.no_grad(): + if kernel == "les": + long_range.kernel.log_alpha.add_(1.0) + else: + long_range.kernel.log_width.sub_(1.0) + after = build(*grid_args, atype=atype_t.reshape(-1))["k_grid"] + torch.testing.assert_close(after.vec, before.vec) + torch.testing.assert_close(after.mask, before.mask) + # the kernel values still follow the live parameters + energy_after = model(coord_t, atype_t, box=box_t)["energy"] + assert not torch.allclose(energy_after, energy_before) + # and the drifted kernel still takes a finite gradient through the graph + model.train() + out = model(coord_t, atype_t, box=box_t) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + kernel_grads = [parameter.grad for parameter in long_range.kernel.parameters()] + assert kernel_grads and all( + grad is not None and torch.isfinite(grad).all() for grad in kernel_grads + ) + + +# --------------------------------------------------------------------------- +# Periodic PME mesh mode +# --------------------------------------------------------------------------- + + +def _periodic_pme_model(kernel: str, **kwargs: object) -> torch.nn.Module: + """A periodic model in the PME mesh mode, ``pme_spacing`` and ``pme_order`` in ``kwargs``.""" + pme_kwargs = { + key: kwargs.pop(key) for key in ("pme_spacing", "pme_order") if key in kwargs + } + config = _config(kernel, **kwargs) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "pme" + config["long_range"].update(pme_kwargs) + return get_model(config).to(env.DEVICE).eval() + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_matches_kspace_reference(kernel: str) -> None: + """With fixed charges the PME energy approaches the converged k-space sum.""" + model = _periodic_pme_model(kernel, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + coord, atype, box = _periodic_case() + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + charge = np.zeros((6, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] + expected = _kspace_reference(coord[0], charge, box[0], kernel) + # The default mesh approximates the reciprocal sum within a few parts in a + # million; the energy, less the self term, is several times smaller. + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-5) + np.testing.assert_allclose(out["atom_energy"].sum().item(), expected, rtol=1e-5) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_force_virial_match_finite_differences(kernel: str) -> None: + """Force and strain virial are the exact derivatives of the PME energy.""" + model = _periodic_pme_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + out = model(coord_t, atype_t, box=box_t) + + def energy(c: np.ndarray, b: np.ndarray) -> float: + return model(_tensor(c), atype_t, box=_tensor(b))["energy"].item() + + h = 1e-5 + force_fd = np.zeros((6, 3)) + for i in range(6): + for d in range(3): + plus, minus = coord.copy(), coord.copy() + plus[0, i, d] += h + minus[0, i, d] -= h + force_fd[i, d] = -(energy(plus, box) - energy(minus, box)) / (2 * h) + np.testing.assert_allclose( + out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-5 + ) + + # the strain family of the mesh strain leaf: r -> (I + eps) r and + # B -> B (I + eps)^T leave the fractional coordinates invariant + virial_fd = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + strain = np.zeros((3, 3)) + strain[a, b] = h + plus = coord @ (np.eye(3) + strain).T + minus = coord @ (np.eye(3) - strain).T + box_plus = box @ (np.eye(3) + strain).T + box_minus = box @ (np.eye(3) - strain).T + virial_fd[a, b] = -(energy(plus, box_plus) - energy(minus, box_minus)) / ( + 2 * h + ) + np.testing.assert_allclose( + out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-4 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_frames_are_independent(kernel: str) -> None: + """Frames on one mesh are independent; a larger cell sharing the batch refines it. + + The frames of a batch share the mesh that resolves the largest cell, so a + smaller cell evaluated beside a larger one sees a finer mesh, which moves + its energy by less than the accuracy of the mesh. + """ + model = _periodic_pme_model(kernel) + coord, atype, box = _periodic_case() + atype2 = np.concatenate([atype, atype], axis=0) + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + for box2, rtol in ( + (np.concatenate([box, box], axis=0), 1e-10), + (np.concatenate([box, box * 1.3], axis=0), 1e-6), + ): + batched = model( + _tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2) + ) + for frame in range(2): + alone = model( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=rtol, atol=1e-10 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_serves_non_periodic_batches(kernel: str) -> None: + """A PME model keeps its non-periodic pair evaluation.""" + model = _periodic_pme_model(kernel) + plain = _model(kernel) + plain.load_state_dict(model.state_dict(), strict=False) + coord, atype = _batch() + mesh_out = model(coord, atype) + plain_out = plain(coord, atype) + torch.testing.assert_close(mesh_out["energy"], plain_out["energy"]) + torch.testing.assert_close(mesh_out["force"], plain_out["force"]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_serialization_round_trip(kernel: str) -> None: + """The record keeps the pme mode and the order, and evaluates identically.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + coord, atype, box = _periodic_case() + model = _periodic_pme_model(kernel, pme_order=6) + record = model.serialize() + assert record["long_range"]["k_grid_mode"] == "pme" + assert record["long_range"]["pme_order"] == 6 + reference = BaseModelDP.deserialize(record) + assert reference.atomic_model.long_range.k_grid_mode == "pme" + expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + ret = reference.call_common(coord, atype, box, neighbor_graph_method="dense") + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-9, + atol=1e-9, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_periodic_pme_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the PME model like the eager one.""" + model = _periodic_pme_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model.eval()(coord_t, atype_t, box=box_t) + model.train() + out = runner(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8) + # the backward reaches the charge network and the kernel through the + # compiled graph + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] + assert grads and all(grad is not None for grad in grads) + assert all(torch.isfinite(grad).all() for grad in grads) diff --git a/source/tests/pt_expt/test_change_bias.py b/source/tests/pt_expt/test_change_bias.py index 719ada14b3..1e042b5fd3 100644 --- a/source/tests/pt_expt/test_change_bias.py +++ b/source/tests/pt_expt/test_change_bias.py @@ -541,6 +541,8 @@ def _make_mock_model(self): fitting_net = MagicMock() class FakeModel(DPModelCommon): + atomic_model = MagicMock(**{"has_long_range.return_value": False}) + def get_out_bias(self): return np.array([[0.0, 0.0]]) diff --git a/source/tests/pt_expt/test_finetune.py b/source/tests/pt_expt/test_finetune.py index 859e1df8c2..87d4fd9f58 100644 --- a/source/tests/pt_expt/test_finetune.py +++ b/source/tests/pt_expt/test_finetune.py @@ -868,6 +868,79 @@ def test_finetune_random_fitting_into_bridged(self) -> None: """A RANDOM fitting of a bridged target takes the plain descriptor.""" self._finetune_across_bridging(source_bridged=False, model_branch="RANDOM") + def test_finetune_short_range_into_long_range(self) -> None: + """A checkpoint without a long-range module fine-tunes into one with it. + + The descriptor and the fitting net come from the checkpoint, while the + new periodic module starts from its initialization and takes the + input statistics of the fine-tuning data, here the bounds of its fixed + k grid, before the bias adjustment evaluates the complete model. + """ + plain = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "seed": 1, + }, + "fitting_net": {"neuron": [16, 16], "seed": 1}, + "data_stat_nbatch": 1, + } + target = deepcopy(plain) + target["long_range"] = { + "type": "latent_charge", + "charge_net": {"neuron": [8], "seed": 1}, + "kernel": {"type": "les", "alpha": 1.0}, + "charge_constraint": {"source": "neutral"}, + "periodic": True, + "k_grid_mode": "fixed", + } + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_long_range_") + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + config = normalize( + update_deepmd_input( + _make_config(self.data_dir, plain, numb_steps=1), warning=False + ) + ) + ckpt_path = self._train_pretrained(config, tmpdir) + ft_config = normalize( + update_deepmd_input( + _make_config(self.data_dir, target, numb_steps=1), warning=False + ) + ) + ft_config["model"], finetune_links = get_finetune_rules( + ckpt_path, ft_config["model"] + ) + trainer_ft = get_trainer( + ft_config, + finetune_model=ckpt_path, + finetune_links=finetune_links, + ) + pretrained_state = torch.load( + ckpt_path, map_location=DEVICE, weights_only=True + )["model"] + ft_state = trainer_ft.wrapper.state_dict() + learned = [ + key + for key in ft_state + if ".descriptor." in key + or (".fitting_net." in key and "bias_atom_e" not in key) + ] + self.assertTrue(learned) + for key in learned: + torch.testing.assert_close(ft_state[key], pretrained_state[key]) + long_range = trainer_ft.wrapper.model["Default"].atomic_model.long_range + self.assertTrue(torch.all(long_range.k_grid_bounds > 0)) + trainer_ft.run() + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + def test_finetune_from_pte(self) -> None: """Train -> freeze to .pte -> finetune from .pte -> verify checkpoint.""" from deepmd.pt_expt.entrypoints.main import ( From 8264c3a363a021288622c26a68681ea002bd3252 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 25 Oct 2026 12:48:22 +0800 Subject: [PATCH 02/15] perf(long-range): accelerate long-range DPA4C inference in LAMMPS Periodic systems take part in the CUDA-graph replay of a single rank. The cell reaches the model in pinned memory of the C++ backend, which the operator reads on the device, so that a replay evaluates the cell of its call; the replay key holds the mesh size of the cell, and after the size changes the model runs eagerly for 64 calls before it captures again, so that a barostat near a boundary of the mesh size does not capture at every crossing. A call with another cell waits for the earlier evaluations before it rewrites that memory, and an evaluation on another stream than the previous one waits on the device for the previous one. The operator sizes the meshes on the host, from the same rounded squares of the lattice vectors as the backend, and forms the geometry of the cells on the device, from pinned cells in place or from a staged copy of pageable ones; a CUDA graph that captures periodic frames needs pinned cells. The cuFFT plans of a mesh shape, which also hold its assignment factors, stay cached while a CUDA graph that runs them lives; a capture of a mesh that no eager call has evaluated is refused with an error naming the mesh, which the C++ backend never meets, since it evaluates every key eagerly before it captures it. A compressed DPA4C model whose long-range module has a charge network of one channel that the fused fitting kernels reproduce, as the defaults give, is evaluated by the operator dpa4c_canonical_compress_long_range_energy_force instead of a composition of operators: one descriptor pass, the charge network, latent_charge_sum on the charges of the owned real nodes, then, in runs of DP_NODE_TILE nodes, the energy fitting and the backward of the charge network seeded with the charge gradient, one descriptor backward with the fixed-point force assembly of both cotangents, and the epilogue adding the energy, the forces at fixed charges and the per-node virials of the long range. The force assembly then reads the source CSR only for the per-atom virial; the archive records this in its metadata (canonical_requires_source_csr), so that LAMMPS builds the source CSR for such a model only at the steps that request the per-atom virial. Archives without the record, other modules and libraries without the operator keep the composition. The direct sum of non-periodic frames runs a system whose blocks of 128 owners would fill less than a third of the resident block slots of the device in blocks and tiles of one warp, in chunks of two tiles counted from the first partner of each block, so that the chunks of a block do not depend on the partners past its frames; blocks of 128 take at most one chunk per two tiles, rounded up, so that no chunk of three tiles sets the duration of a launch that fits in one wave. The influence function of a frame takes no more blocks than the rows of its half spectrum fill and reads its wave-vector factors and scale from the geometry of the frame, and the float64 kernels of the node stencils and of the owners spread their warps over the multiprocessors. Both tile widths of the pair kernel are compiled, which lengthens the compilation of direct_sum.cu from 58 to 74 s and adds 1.9 MB to the operator library. LAMMPS with Kokkos on compute capability 12.0, trained Air model with an LES or SOG module, ms per step, against the libraries and frozen models of the long-range module as introduced: 512 atoms, periodic LES 0.665 -> 0.280 and SOG 0.644 -> 0.277, cluster LES 0.343 -> 0.264 and SOG 0.353 -> 0.263; 10,296 atoms, periodic LES 1.469 -> 1.242 and SOG 1.357 -> 1.137, cluster LES 1.274 -> 1.166 and SOG 1.479 -> 1.302; 100,800 atoms, in 1000-step runs on two GPUs, periodic LES 11.32 -> 10.38 and 11.42 -> 10.11 and SOG 10.03 -> 8.98 and 9.81 -> 8.93, cluster LES 16.22 -> 14.76 and 16.00 -> 14.56 and SOG 32.13 -> 31.63 and 32.33 -> 31.52. A periodic LES system of 512 atoms whose cell is deformed at every step, every 2 or every 20 steps replays in 0.28 to 0.29 ms per step, as with a constant cell, against 0.35 to 0.38 ms evaluated eagerly by the same library. Replayed and eager trajectories agree bit for bit over 200 steps of NVE and NPT. Against the module as introduced, the step-0 forces of the eight systems move by up to 3.2e-7 of the largest and their energies by up to 1.0e-6 eV, the rounding of the fused evaluation and of the float32 pair sums. Short-range models keep their kernel time (a step of the cluster of 100,800 atoms 7.24 and 7.21 -> 7.20 and 7.19 ms on one GPU); their step times stay within 2.5%, except on that cluster, whose step time varies by up to 20% from run to run. --- deepmd/pt_expt/kernels/dpa4c/canonical.py | 37 + .../kernels/dpa4c/canonical_long_range.py | 174 +++++ deepmd/pt_expt/kernels/long_range.py | 7 +- deepmd/pt_expt/long_range/latent_charge.py | 33 +- deepmd/pt_expt/model/ener_model.py | 88 ++- deepmd/pt_expt/utils/serialization.py | 29 + doc/model/long-range.md | 13 +- source/api_c/include/c_api.h | 30 +- source/api_cc/include/DeepPot.h | 44 +- source/api_cc/include/DeepPotPTExpt.h | 28 +- source/api_cc/src/DeepPotPTExpt.cc | 111 ++- source/api_cc/src/commonPTExpt.h | 23 +- source/api_cc/src/longRangePTExpt.h | 158 ++++- .../test_deeppot_dpa4c_long_range_ptexpt.cc | 321 ++++++++- source/op/pt/dpa4c/graph_compress.cu | 652 +++++++++++++++--- source/op/pt/dpa4c/ops.cc | 32 + source/op/pt/long_range/charge_sum.h | 17 + source/op/pt/long_range/direct_sum.cu | 203 ++++-- source/op/pt/long_range/latent_charge_sum.cu | 11 +- source/op/pt/long_range/mesh_sum.cu | 392 ++++++++--- source/op/pt/long_range/ops.cc | 14 +- source/tests/infer/gen_dpa4c_long_range.py | 3 + .../pt_expt/kernels/test_long_range_sum.py | 112 ++- source/tests/pt_expt/model/test_long_range.py | 217 ++++++ 24 files changed, 2311 insertions(+), 438 deletions(-) diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical.py b/deepmd/pt_expt/kernels/dpa4c/canonical.py index 8fed0b886e..4c3f8eedd8 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical.py @@ -96,6 +96,13 @@ def op_available() -> bool: ) +def long_range_op_available() -> bool: + """Return whether the backend device carries the fused long-range operator.""" + return op_available() and operator_available( + "dpa4c_canonical_compress_long_range_energy_force" + ) + + #: Position of the per-node magnetic moment among the compression arguments, #: which open with the radial table (see ``compressed_operator_arguments``), #: and among the inputs that follow ``edge_vec`` in the compact descriptor @@ -246,6 +253,32 @@ def _energy_force_fake( ) +def _long_range_energy_force_fake( + edge_vec: torch.Tensor, + source: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_row_ptr: torch.Tensor, + source_order: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + atype: torch.Tensor, + *args: Any, +) -> tuple[torch.Tensor, ...]: + del source, destination_row_ptr, source_row_ptr, source_order, n_local, args + nodes = atype.shape[0] + frames = n_node.shape[0] + # Each output is allocated separately, because the schema declares six + # unannotated results and no two of them may share storage. + return ( + edge_vec.new_empty(frames, 1, dtype=torch.float64), + edge_vec.new_empty(nodes, 1, dtype=torch.float64), + edge_vec.new_empty(nodes, 3, dtype=torch.float64), + edge_vec.new_empty(frames, 3, 3, dtype=torch.float64), + edge_vec.new_empty(nodes, 3, 3, dtype=torch.float32), + edge_vec.new_empty(nodes, dtype=torch.int32), + ) + + def _cpu_energy_force( edge_vec: torch.Tensor, source: torch.Tensor, @@ -537,6 +570,10 @@ def ensure_registered() -> None: torch.library.register_fake("deepmd::dpa4c_canonical_compress_energy_force")( _energy_force_fake ) + if long_range_op_available(): + torch.library.register_fake( + "deepmd::dpa4c_canonical_compress_long_range_energy_force" + )(_long_range_energy_force_fake) torch.library.register_fake("deepmd::dpa4c_canonical_compress_backward_inplace")( _backward_inplace_fake ) diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py index 3bb786ff58..e7bd81b80f 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -18,6 +18,180 @@ ) +def fused_long_range_eligible(long_range: Any) -> bool: + """Whether the fused long-range operator evaluates a long-range module. + + The operator evaluates a charge network of one channel that the fused + fitting kernels reproduce, beside the compressed descriptor and the energy + fitting that the compact canonical lower already requires; a module whose + charge network they do not reproduce, or a library without the operator, + takes :func:`dpa4c_canonical_long_range_energy_force`. + """ + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + long_range_op_available, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_eligible, + ) + + return long_range_op_available() and fitting_eligible(long_range.charge_net) + + +def dpa4c_canonical_compress_long_range_energy_force( + descriptor: Any, + fitting: Any, + long_range: Any, + graph: CanonicalGraph, + atype: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor | None, + box: torch.Tensor | None, + atom_bias: torch.Tensor, + fitting_routes: list[int] | None = None, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Evaluate compressed DPA4C and its long-range module in one operator. + + The operator evaluates the descriptor of every node, the charges of the + module from it, the energy fitting, the long range of the charges and one + descriptor backward from the cotangents of both networks, and assembles + the forces in fixed point like + :func:`~deepmd.pt_expt.kernels.dpa4c.canonical.dpa4c_canonical_compress_energy_force`, + so the force assembly reads the source CSR only for the per-atom virial. + It derives the ownership mask from the frame counts, and the charged + nodes of the module are the owned real nodes. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor without a spin condition. + fitting + Eligible pt_expt energy fitting network. + long_range + Long-range module for which :func:`fused_long_range_eligible` holds. + graph + Canonical graph as for + :func:`~deepmd.pt_expt.kernels.dpa4c.canonical.dpa4c_canonical_compress_energy_force`; + its source CSR is either complete or a two-entry placeholder, which + declines the per-atom virial. + atype + Flat node atom types with shape ``(N,)``; a negative type marks a + virtual node, which contributes no energy and carries no charge. + coord + Node positions with shape ``(N, 3)`` in Å. + communicator + The address of the MPI_Comm handle of a frame decomposed into + domains, as an int64 CPU tensor of one element, or zero or ``None`` + for a single domain. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the CPU, a + zero row marking a complete non-periodic frame, or ``None`` when no + frame is periodic. + atom_bias + Combined atomic energy bias with shape ``(ntypes,)`` in eV. + fitting_routes + Fitting-route table of the energy fitting, as for the fused operator + without a long-range module, or ``None`` for the measured table. + + Returns + ------- + energy + Per-frame energy with shape ``(F, 1)`` in eV, fp64. + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, fp64. + force + Per-node force with shape ``(N, 3)`` in eV/Å, fp64. + virial + Per-frame virial with shape ``(F, 3, 3)`` in eV, fp64. + atom_virial + Per-node virial with shape ``(N, 3, 3)`` in eV, fp32; zero when the + graph carries no source CSR. + mask + Ownership mask with shape ``(N,)``, int32. + """ + from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, + ) + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + ensure_registered, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + compressed_operator_arguments, + ) + from deepmd.pt_expt.kernels.fitting_routes import fitting_routes as measured_routes + from deepmd.pt_expt.kernels.graph_fitting import ( + ensure_registered as ensure_fitting_registered, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_operator_arguments, + node_tile, + ) + from deepmd.pt_expt.kernels.long_range import ( + kernel_operator_arguments, + ) + + ensure_registered() + ensure_fitting_registered() + network = fitting_operator_arguments(fitting) + charge_net = long_range.charge_net + charges = fitting_operator_arguments(charge_net) + total_charge, channel = long_range.operator_constraint(graph.n_node.shape[0], coord) + kernel, coefficient = kernel_operator_arguments(long_range.kernel) + if communicator is None: + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") + routes = measured_routes(fitting) if fitting_routes is None else fitting_routes + return torch.ops.deepmd.dpa4c_canonical_compress_long_range_energy_force( + graph.edge_vec, + graph.source, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + graph.n_node, + graph.n_local, + atype, + *compressed_operator_arguments(descriptor), + int(descriptor.lmax), + *descriptor._compression_scalars, + network.weights, + network.biases, + network.residuals, + network.head_weight, + network.head_bias, + atom_bias.to(torch.float64).contiguous(), + network.activation, + node_tile(), + charges.weights, + charges.biases, + charges.residuals, + charges.head_weight, + charges.head_bias, + charge_net.bias_atom_e[:, 0].to(torch.float64).contiguous(), + charges.activation, + coord.to(torch.float64).contiguous(), + total_charge, + # A module that does not serve periodic systems sums every frame + # directly. + box if long_range.periodic else None, + coefficient.to(coord.device), + communicator, + kernel, + channel, + COULOMB_CONSTANT, + long_range.pme_order, + long_range.mesh_spacing() if long_range.periodic else 1.0, + # No table is passed as None rather than as an empty list: the + # AOTInductor proxy executor of torch 2.13 and earlier hands a + # serialized empty int list to the operator as None. + list(routes) if routes else None, + ) + + def dpa4c_canonical_long_range_energy_force( descriptor: Any, fitting: Any, diff --git a/deepmd/pt_expt/kernels/long_range.py b/deepmd/pt_expt/kernels/long_range.py index 89db40ad93..33ab82a828 100644 --- a/deepmd/pt_expt/kernels/long_range.py +++ b/deepmd/pt_expt/kernels/long_range.py @@ -430,7 +430,12 @@ def latent_charge_sum( CPU, the rows of every row-major 3x3 cell being its lattice vectors, or ``None`` for non-periodic frames; a zero row marks a non-periodic frame. A periodic frame is evaluated on a mesh of spacing at most - ``spacing``, with B-splines of order ``order``. + ``spacing``, with B-splines of order ``order``. The device reads cells + in pinned memory in place while the stream runs the call, so the + caller keeps them alive and unchanged until then. A CUDA graph that + captures the call needs them pinned and reads them at every replay, + and it needs an eager call to have evaluated meshes of the same sizes + with the same channel count and ``order``. order The order of the B-splines of the mesh, at most :data:`MAX_ORDER`. spacing diff --git a/deepmd/pt_expt/long_range/latent_charge.py b/deepmd/pt_expt/long_range/latent_charge.py index b79dffef4c..cadd860d50 100644 --- a/deepmd/pt_expt/long_range/latent_charge.py +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -76,6 +76,31 @@ def operator_eligible(self) -> bool: and not self.uses_charge_spin() ) + def operator_constraint( + self, n_frames: int, like: torch.Tensor + ) -> tuple[torch.Tensor, int]: + """The total charge of every frame and the constrained channel of the operator. + + Parameters + ---------- + n_frames + The number of frames. + like + A tensor giving the device of the total charges. + + Returns + ------- + tuple[torch.Tensor, int] + The total charges with shape ``(n_frames,)``, float64 and + contiguous, and the constrained channel, or -1 without a + constraint, in which case the total charges are zero. + """ + total_charge = like.new_zeros(n_frames, dtype=torch.float64) + if self.charge_constraint is None: + return total_charge, -1 + total_charge = self._total_charge(n_frames, total_charge, None, None) + return total_charge.contiguous(), int(self.charge_constraint["channel"]) + def energy_force_from_coord( self, descriptor: torch.Tensor, @@ -144,17 +169,13 @@ def energy_force_from_coord( leaf = descriptor.detach().requires_grad_(True) with torch.enable_grad(): charge = self.charge_net.call_graph(leaf, atype)["latent_charge"] - total_charge = coord.new_zeros(n_node.shape[0], dtype=torch.float64) - channel = -1 - if self.charge_constraint is not None: - channel = self.charge_constraint["channel"] - total_charge = self._total_charge(n_node.shape[0], total_charge, None, None) + total_charge, channel = self.operator_constraint(n_node.shape[0], coord) atom_energy, charge_gradient, force, atom_virial, virial = latent_charge_sum( coord.to(torch.float64).contiguous(), charge.detach().to(torch.float32).contiguous(), node_mask.contiguous(), n_node, - total_charge.contiguous(), + total_charge, self.kernel, channel, COULOMB_CONSTANT, diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 61c7bf103a..f5f73547ac 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -174,8 +174,9 @@ def forward_lower_canonical_graph( Destination CSR offsets with shape ``(N + 1,)``, int64. source_row_ptr Source CSR offsets with shape ``(N + 1,)``, int64. A DPA4C model - also accepts a two-entry placeholder, which declines the - per-atom virial; that output is then zero. + evaluated by a fused operator, which every DPA4C model without a + long-range module is, also accepts a two-entry placeholder, which + declines the per-atom virial; that output is then zero. source_order Source-grouped edge positions with shape ``(S,)``, int64, or a two-entry placeholder alongside a placeholder ``source_row_ptr``. @@ -239,9 +240,24 @@ def forward_lower_canonical_graph( descriptor = learned.descriptor fitting = learned.fitting_net atom_bias = fused_atom_bias(self.atomic_model, learned) + + def fused_graph() -> CanonicalGraph: + # A program without the per-atom virial declines it with the + # two-entry placeholders of the source CSR, so that a fused + # operator stores no edge cotangent whatever its caller passes. + if do_atomic_virial: + return graph + return dataclasses.replace( + graph, + source_row_ptr=source_row_ptr.new_zeros(2), + source_order=source_order.new_zeros(2), + ) + if learned.has_long_range(): from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + dpa4c_canonical_compress_long_range_energy_force, dpa4c_canonical_long_range_energy_force, + fused_long_range_eligible, ) if coord is None: @@ -249,39 +265,47 @@ def forward_lower_canonical_graph( "a model with a long-range module reads the node positions" ) check_owned_frames(n_node, n_local, communicator, box) - atype, output_mask = self.atomic_model._prepare_graph_nodes( - n_node, - n_local, - atype, - edge_vec, - ) force_mag = None - mask = output_mask.to(torch.int32) - energy, atom_energy, force, virial, atom_virial = ( - dpa4c_canonical_long_range_energy_force( - descriptor, - fitting, - learned.long_range, - graph, + if fused_long_range_eligible(learned.long_range): + # The fused operator derives the ownership mask from the frame + # counts itself, as the operator without a long-range module. + energy, atom_energy, force, virial, atom_virial, mask = ( + dpa4c_canonical_compress_long_range_energy_force( + descriptor, + fitting, + learned.long_range, + fused_graph(), + atype, + coord, + communicator, + box, + atom_bias, + ) + ) + else: + atype, output_mask = self.atomic_model._prepare_graph_nodes( + n_node, + n_local, atype, - coord, - communicator, - box, - output_mask, - atom_bias, - do_atomic_virial, + edge_vec, ) - ) - elif use_dpa4c: - # A program without the per-atom virial declines it with the - # two-entry placeholders of the source CSR, so that the fused - # operator stores no edge cotangent whatever its caller passes. - if not do_atomic_virial: - graph = dataclasses.replace( - graph, - source_row_ptr=source_row_ptr.new_zeros(2), - source_order=source_order.new_zeros(2), + mask = output_mask.to(torch.int32) + energy, atom_energy, force, virial, atom_virial = ( + dpa4c_canonical_long_range_energy_force( + descriptor, + fitting, + learned.long_range, + graph, + atype, + coord, + communicator, + box, + output_mask, + atom_bias, + do_atomic_virial, + ) ) + elif use_dpa4c: # The fused operator derives the ownership mask from the frame # counts itself, so nothing runs on the node axis around it. The # fitting-route table measured for the network on the backend @@ -290,7 +314,7 @@ def forward_lower_canonical_graph( dpa4c_canonical_compress_energy_force( descriptor, fitting, - graph, + fused_graph(), atype, atom_bias, spin, diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index f7d9ad2511..16172cf7b3 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -149,6 +149,32 @@ def _graph_reads_source_csr(exported: Any) -> bool: return not found +def _canonical_requires_source_csr(model: Any, lower_kind: str) -> bool: + """Return whether a compact canonical lower reads the source CSR on every call. + + The force assembly of compressed DPA1 walks the source CSR, and so does + the composition of operators that evaluates a long-range module whose + charge network the fused long-range operator does not reproduce. The + fused DPA4C operators assemble the forces in fixed point and read the + source CSR only for the per-atom virial, so a caller that declines it may + pass two-entry placeholders in its place. + """ + if lower_kind == "dpa1_canonical": + return True + decomposition = model.atomic_model.fused_decomposition() + learned = None if decomposition is None else decomposition[0] + if learned is None: + # The compact lower serves no such model, which the export reports. + return True + if not learned.has_long_range(): + return False + from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + fused_long_range_eligible, + ) + + return not fused_long_range_eligible(learned.long_range) + + def _numpy_to_json_serializable(model_obj: dict) -> dict: """Convert numpy arrays in a model dict to JSON-serializable lists.""" return traverse_model_dict( @@ -1446,6 +1472,9 @@ def _probe_has_message_passing(obj: object) -> bool | None: if lower_kind in ("dpa1_canonical", "dpa4c_canonical"): meta["canonical_index_dtype"] = "auto" meta["canonical_index_storage"] = "byte_words" + meta["canonical_requires_source_csr"] = _canonical_requires_source_csr( + model, lower_kind + ) # Model-level pair-type exclusion (``pair_exclude_types``): a list of # ``[ti, tj]`` type pairs whose interaction is dropped. Exclusion is a diff --git a/doc/model/long-range.md b/doc/model/long-range.md index 4c22bc5f96..b883b463fa 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -236,10 +236,15 @@ memory grows linearly with the number of atoms and its cost with their square, and evaluates a periodic system on the mesh of its cell. The operator serves the `les` kernel and the `sog` kernel of up to 32 Gaussians, with up to 8 charge channels, B-splines of order up to 12, a charge network without -frame or atomic parameters, and no `pair_exclude_types`. Any other model, -like an uncompressed one, receives the pairs or the mesh of each frame from -the interface, so the memory of a non-periodic evaluation grows with the -square of the number of atoms. +frame or atomic parameters, and no `pair_exclude_types`. A model the +operator does not serve, like an uncompressed one, receives the pairs or the +mesh of each frame from the interface, so the memory of a non-periodic +evaluation grows with the square of the number of atoms. The operator +evaluates the long-range term together with the descriptor and both networks, +which is faster, when the module has one charge channel and a charge network +in `float32` with `tanh` or `silu` hidden layers whose widths are multiples of +four and without `resnet_dt`; the defaults give such a module in a compressed +model. The following are not supported for a model with a long-range module: multi-rank (domain-decomposed) inference of a model that the dedicated diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 2bfe93694c..ff8903d341 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -480,18 +480,24 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, * without a synchronization of the stream. A graph without halo nodes is * captured into a CUDA graph on the second consecutive call with the same * buffers and extents and replayed while they stay the same, except where a - * model with a long-range module receives a communicator or a cell. The host - * waits for the device when a call captures or releases a graph, which - * happens on that second call and on a call whose buffers or extents change. - * DP_DeepPotSetChargeSpin waits for the work in flight on the model device - * before it replaces the constants and for the replacement before it returns, - * and the caller completes the work on its stream before it deletes the model. - * NULL runs on the current stream of the calling thread for the model device - * and waits for that stream before returning. Models in different threads of - * one process evaluate concurrently, and their captures, releases of graphs - * and charge-state updates, which wait for the whole device, take turns; CUDA - * work outside DeePMD-kit that waits for the whole device while a capture is - * under way invalidates it, and the model then evaluates without replay. + * model with a long-range module receives a communicator. A model whose module + * serves periodic systems copies the cell into memory of its own, which a + * replay reads, and the mesh size of the cell counts as an extent: a call with + * another cell than the previous call first waits for the earlier evaluations, + * consecutive evaluations on different streams run in call order, and after the + * mesh size changes, the model evaluates without replay for 64 calls, after + * which the capture follows as for new extents. The host waits for the device + * when a call captures or releases a graph, which happens on that second call + * and on a call whose buffers or extents change. DP_DeepPotSetChargeSpin waits + * for the work in flight on the model device before it replaces the constants + * and for the replacement before it returns, and the caller completes the work + * on its stream before it deletes the model. NULL runs on the current stream of + * the calling thread for the model device and waits for that stream before + * returning. Models in different threads of one process evaluate concurrently, + * and their captures, releases of graphs and charge-state updates, which wait + * for the whole device, take turns; CUDA work outside DeePMD-kit that waits for + * the whole device while a capture is under way invalidates it, and the model + * then evaluates without replay. * @note API versions 28 to 30 named this function * DP_DeepPotComputeCanonicalGraphGPU, without ``d_virial``, ``d_coord``, * ``world``, ``box`` and ``stream``; a binary built against that name diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index a7fcdc6479..53f7091e79 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -967,25 +967,31 @@ class DeepPot : public DeepBaseModel { * @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. - * @param[in] stream Accelerator stream (``cudaStream_t``) of the model - * device on which the inputs were produced and the outputs will be - * consumed; a non-null stream requires a CUDA build. The evaluation is - * enqueued on it without a synchronization of the stream. A graph without - * halo nodes is captured into a CUDA graph on the second consecutive call - * with the same buffers and extents and replayed while they stay the same, - * except where a model with a long-range module receives a communicator or - * a cell. The host waits for the device when a call captures or releases a - * graph, which happens on that second call and on a call whose buffers or - * extents change. set_charge_spin waits for the work in flight on the model - * device before it replaces the constants and for the replacement before it - * returns, and the caller completes the work on its stream before it - * destroys the model. A null stream runs on the current stream of the - * calling thread for the model device and waits for that stream before - * returning. Models in different threads of one process evaluate - * concurrently, and their captures, releases of graphs and charge-state - * updates, which wait for the whole device, take turns; CUDA work outside - * DeePMD-kit that waits for the whole device while a capture is under way - * invalidates it, and the model then evaluates without replay. + * @param[in] stream Accelerator stream (``cudaStream_t``) of the model device + * on which the inputs were produced and the outputs will be consumed; a + * non-null stream requires a CUDA build. The evaluation is enqueued on it + * without a synchronization of the stream. A graph without halo nodes is + * captured into a CUDA graph on the second consecutive call with the same + * buffers and extents and replayed while they stay the same, except where a + * model with a long-range module receives a communicator. A model whose + * module serves periodic systems copies the cell into memory of its own, + * which a replay reads, and the mesh size of the cell counts as an extent: a + * call with another cell than the previous call first waits for the earlier + * evaluations, consecutive evaluations on different streams run in call + * order, and after the mesh size changes, the model evaluates without replay + * for 64 calls, after which the capture follows as for new extents. The host + * waits for the device when a call captures or releases a graph, which + * happens on that second call and on a call whose buffers or extents change. + * set_charge_spin waits for the work in flight on the model device before it + * replaces the constants and for the replacement before it returns, and the + * caller completes the work on its stream before it destroys the model. A + * null stream runs on the current stream of the calling thread for the model + * device and waits for that stream before returning. Models in different + * threads of one process evaluate concurrently, and their captures, releases + * of graphs and charge-state updates, which wait for the whole device, take + * turns; CUDA work outside DeePMD-kit that waits for the whole device while a + * capture is under way invalidates it, and the model then evaluates without + * replay. */ void compute_canonical_graph_gpu(double* d_atom_energy, double* d_force, diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 71e66a2109..bc17f23c01 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -14,17 +14,21 @@ #include +#include +#include #include +#include #include "DeepPot.h" #include "graph_assembly.h" // Forward-declare to keep these out of the public header. Defined in -// commonPTExpt.h. +// commonPTExpt.h, and PinnedCell in longRangePTExpt.h. namespace deepmd::ptexpt { class TempFile; class ChargeStateFold; class GraphReplay; +class PinnedCell; } // namespace deepmd::ptexpt namespace torch::inductor { @@ -513,9 +517,18 @@ 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; + // The cell input of the device-resident evaluations of a model with a + // long-range module, in pinned memory that a replay reads again. + std::unique_ptr long_range_cell_; + // The mesh size of the cell of the last device-resident evaluation of a + // model with a long-range module, zero for a non-periodic frame, and the + // evaluations since it last changed, counted up to the number of calls for + // which a changed mesh is evaluated eagerly before it is captured again. + std::optional > long_range_mesh_; + std::int64_t long_range_mesh_calls_ = 0; // CUDA-graph replay of repeated device-resident canonical evaluations, - // declared after the loader so that the graph is released before the - // constants it reads. + // declared after the loader and the cell so that the graph is released + // before the constants and the cell it reads. std::unique_ptr canonical_replay_; // Serializes the runs of the model, which runs on a single-threaded runner, // the device-resident entry points, whose replay of canonical evaluations @@ -688,9 +701,10 @@ class DeepPotPTExpt : public DeepPotBackend { * @param[in] world The MPI communicator (a pointer to an ``MPI_Comm``) of the * ranks a decomposed system spans, or NULL for a system that one rank holds * whole; a model with a long-range module receives it after the positions. - * @param[in] box The row-major cell of a periodic frame, nine values, or - * empty for a non-periodic one; a model with a long-range module receives - * it after the communicator. + * @param[in] cell The cell, a float64 CPU tensor of shape ``(1, 9)``: the + * row-major cell of a periodic frame, or zero for a non-periodic one; a + * model with a long-range module receives it after the communicator, and + * every other model ignores it. */ std::vector run_model_canonical_graph( const torch::Tensor& atype, @@ -703,7 +717,7 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& source_order, const torch::Tensor& coord, const void* world, - const std::vector& box); + const torch::Tensor& cell); /** * @brief Evaluate the long-range module on a rank that holds no node. diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index e6ec9d9dd5..dba250c575 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -131,6 +131,15 @@ at::Tensor communicator_tensor(const void* world) { torch::TensorOptions().dtype(torch::kInt64)); } +// A capture waits for the device, empties the cache of the allocator and +// allocates the memory of its graph anew, which takes far longer than an +// evaluation of a small system. After the mesh of its cell changes size, a +// model with a long-range module is therefore evaluated eagerly until the +// size has held for this many calls, so that a cell that fluctuates across a +// boundary of the mesh size under a barostat is not captured at every +// crossing. +constexpr std::int64_t kStableMeshCalls = 64; + // The cell input of the compact lower of a long-range model: a float64 CPU // tensor of shape (1, 9), the rows of the row-major cell its lattice vectors, // zero for a non-periodic frame. @@ -284,7 +293,13 @@ void DeepPotPTExpt::init(const std::string& model, lower_input_is_graph_ = lower_input_kind == "graph"; lower_input_is_canonical_ = lower_input_kind == "dpa1_canonical" || lower_input_kind == "dpa4c_canonical"; - canonical_requires_source_csr_ = lower_input_kind == "dpa1_canonical"; + // Whether the force assembly of the compact lower reads the source CSR on + // every call, as the archive records it; an archive without the record + // reads it for compressed DPA1 and for a long-range module (below). + canonical_requires_source_csr_ = + metadata.obj_val.count("canonical_requires_source_csr") + ? metadata["canonical_requires_source_csr"].as_bool() + : lower_input_kind == "dpa1_canonical"; } else { lower_input_is_edge_ = false; lower_input_is_graph_ = false; @@ -315,8 +330,9 @@ void DeepPotPTExpt::init(const std::string& model, "a model with a long-range module requires the graph or the " "compact canonical lower."); } - // Its compact lower assembles the force through the source CSR. - canonical_requires_source_csr_ = lower_input_is_canonical_; + if (!metadata.obj_val.count("canonical_requires_source_csr")) { + canonical_requires_source_csr_ = lower_input_is_canonical_; + } } if (lower_input_is_edge_) { std::cerr << "WARNING: This .pt2 uses the deprecated edge_vec lower " @@ -649,7 +665,7 @@ std::vector DeepPotPTExpt::run_model_canonical_graph( const torch::Tensor& source_order, const torch::Tensor& coord, const void* world, - const std::vector& box) { + const torch::Tensor& cell) { std::vector inputs = { atype, n_node, n_local, canonicalIndexBytes(source), @@ -658,7 +674,7 @@ std::vector DeepPotPTExpt::run_model_canonical_graph( if (long_range_) { inputs.push_back(coord); inputs.push_back(communicator_tensor(world)); - inputs.push_back(cell_tensor(box)); + inputs.push_back(cell); } const std::lock_guard evaluation(evaluation_mutex_); return loader->run(inputs); @@ -710,7 +726,8 @@ void DeepPotPTExpt::join_long_range_exchange(const void* world, torch::zeros({2}, i64.dtype(torch::kInt32)), torch::zeros({2, 3}, i64.dtype(torch::kFloat32)), torch::zeros({2}, i64), torch::zeros({2}, i64), torch::arange(2, i64), - torch::zeros({1, 3}, i64.dtype(torch::kFloat64)), world, box); + torch::zeros({1, 3}, i64.dtype(torch::kFloat64)), world, + cell_tensor(box)); } std::vector DeepPotPTExpt::run_model_with_comm( @@ -1483,7 +1500,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, compact.atype, compact.n_node, compact.n_local, compact.source, compact.edge_vec, compact.destination_row_ptr, compact.source_row_ptr, compact.source_order, node_coord, world, - long_range_box); + long_range_ ? cell_tensor(long_range_box) : torch::Tensor()); } else { flat_outputs = run_model_graph( node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index, @@ -1871,7 +1888,8 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, 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, node_coord, /*world=*/nullptr, long_range_box); + compact.source_order, node_coord, /*world=*/nullptr, + long_range_ ? cell_tensor(long_range_box) : torch::Tensor()); } else { flat_outputs = run_model_graph( graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, @@ -2871,6 +2889,15 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( has_source_csr ? d_source_row_ptr : source_row_placeholder_.data_ptr(); const void* const source_order_data = has_source_csr ? d_source_order : source_order_placeholder_.data_ptr(); + // A model with a long-range module takes the cell in pinned memory, which + // the evaluation of a periodic system reads on the device. + const bool long_range = has_long_range(); + if (long_range && !long_range_cell_) { + long_range_cell_ = std::make_unique(); + } + const torch::Tensor cell = + long_range ? long_range_cell_->hold(long_range_box, gpu_id) + : torch::Tensor(); // The tensors over the caller's buffers are built inside the evaluation, // which a replay skips: wrapping a device buffer queries the driver for @@ -2911,7 +2938,7 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( extract_outputs(output, run_model_canonical_graph( atype, n_node, n_local, source, edge_vec, destination_row_ptr, source_row_ptr, - source_order, coord, world, long_range_box)); + source_order, coord, world, cell)); auto atom_energy = output["atom_energy"] .reshape({nall_nodes}) .slice(0, 0, nloc) @@ -2944,27 +2971,61 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( // neighbor rebuild, so a graph with halo nodes is evaluated eagerly. A // model with a long-range module reads the node positions, whose address // is then part of the key, and is evaluated eagerly over several ranks, - // where it may communicate between its kernels, and in a cell, which the - // evaluation reads on the host. - const bool long_range = has_long_range(); - if (nloc != nall_nodes || - (long_range && (world != nullptr || box != nullptr))) { + // where it may communicate between its kernels. A replay reads the cell + // of its call, while the mesh size of the cell fixes the buffers and the + // launches of the evaluation, so the size joins the key. + bool replay = nloc == nall_nodes && !(long_range && world != nullptr); + std::array mesh = {}; + if (long_range) { + if (!long_range_box.empty()) { + mesh = ptexpt::meshDims(long_range_box, mesh_spacing_); + } + if (long_range_mesh_ != mesh) { + // The first mesh of the instance is replayed once its key repeats. The + // graph of another mesh is released at once, so that its memory is + // free during the eager calls that follow. + long_range_mesh_calls_ = + long_range_mesh_.has_value() ? 0 : kStableMeshCalls; + long_range_mesh_ = mesh; + if (canonical_replay_) { + canonical_replay_->invalidate(); + } + } + if (long_range_mesh_calls_ < kStableMeshCalls) { + ++long_range_mesh_calls_; + replay = false; + } + } + if (!replay) { evaluate(); } else { if (!canonical_replay_) { canonical_replay_ = std::make_unique(); } - canonical_replay_->run( - {ptexpt::address_key(d_atom_energy), ptexpt::address_key(d_force), - ptexpt::address_key(d_virial), ptexpt::address_key(d_atom_virial), - ptexpt::address_key(d_atype), ptexpt::address_key(d_source), - ptexpt::address_key(d_edge_vec), - ptexpt::address_key(d_destination_row_ptr), - ptexpt::address_key(source_row_data), - ptexpt::address_key(source_order_data), - ptexpt::address_key(long_range ? d_coord : nullptr), nloc, - nall_nodes, edge_storage, index_width}, - gpu_id, evaluate); + ptexpt::GraphReplay::Key key = { + ptexpt::address_key(d_atom_energy), + ptexpt::address_key(d_force), + ptexpt::address_key(d_virial), + ptexpt::address_key(d_atom_virial), + ptexpt::address_key(d_atype), + ptexpt::address_key(d_source), + ptexpt::address_key(d_edge_vec), + ptexpt::address_key(d_destination_row_ptr), + ptexpt::address_key(source_row_data), + ptexpt::address_key(source_order_data), + ptexpt::address_key(long_range ? d_coord : nullptr), + nloc, + nall_nodes, + edge_storage, + index_width}; + if (long_range) { + key.push_back(ptexpt::address_key(cell.data_ptr())); + key.insert(key.end(), mesh.begin(), mesh.end()); + } + canonical_replay_->run(key, gpu_id, evaluate); + } + if (long_range_periodic_) { + long_range_cell_->record(); } if (stream == nullptr) { synchronize_current_accelerator_stream(gpu_id); diff --git a/source/api_cc/src/commonPTExpt.h b/source/api_cc/src/commonPTExpt.h index 4ed044ff26..1981531ccf 100644 --- a/source/api_cc/src/commonPTExpt.h +++ b/source/api_cc/src/commonPTExpt.h @@ -165,19 +165,21 @@ inline void copy_results(const std::vector& targets, * @brief Replay repeated device-resident evaluations as one CUDA graph. * * Launching kernels dominates an evaluation of a small system. Two - * evaluations over the same buffers and extents launch the same kernels with - * the same arguments, so the evaluation is captured into a CUDA graph on the - * second consecutive call with the same key -- every buffer address and - * extent the evaluation depends on -- and relaunched as a whole while the key - * stays the same. A call with another key releases the graph and evaluates - * eagerly. + * evaluations over the same buffers, extents and host values launch the same + * kernels with the same arguments, so the evaluation is captured into a CUDA + * graph on the second consecutive call with the same key -- every buffer + * address, extent and host value that shapes its launches -- and relaunched + * as a whole while the key stays the same. A call with another key releases + * the graph and evaluates eagerly. * * The evaluation must be free of host synchronization, and every device * tensor it uses beyond the caller's buffers must either be allocated inside * it, so that the graph owns them, or persist at an address that the key - * covers. The constants of the model are read where they reside at capture, - * so whoever replaces them calls ``invalidate`` first. An evaluation that - * cannot be captured runs eagerly instead (see ``capture``). + * covers; a replay repeats the launches of the capture without the host code + * that made them, so every host value that shapes them belongs to the key. + * The constants of the model are read where they reside at capture, so + * whoever replaces them calls ``invalidate`` first. An evaluation that cannot + * be captured runs eagerly instead (see ``capture``). * * A capture, like the release of a graph, waits for the whole device, and * both hold ``device_wide_mutex``, so the replays of different models capture @@ -193,7 +195,8 @@ class GraphReplay { /** * @brief Evaluate under ``key`` on the current stream of ``device``. * - * @param key Every buffer address and extent ``evaluate`` depends on. + * @param key Every buffer address, extent and host value that shapes the + * launches of ``evaluate``. * @param device Index of the CUDA device. * @param evaluate The evaluation, enqueued on the current stream. */ diff --git a/source/api_cc/src/longRangePTExpt.h b/source/api_cc/src/longRangePTExpt.h index 352f85a643..f955f1f6f9 100644 --- a/source/api_cc/src/longRangePTExpt.h +++ b/source/api_cc/src/longRangePTExpt.h @@ -3,13 +3,23 @@ #include +#include #include #include #include +#include +#include #include #include "commonPT.h" +#ifdef DEEPMD_TORCH_HAS_CUDA +#include +#include +#include +#include +#endif + namespace deepmd { namespace ptexpt { @@ -93,17 +103,53 @@ inline std::int64_t evenFftSize(std::int64_t n) { } } +/** + * @brief The size of the mesh of a periodic frame along its lattice vectors. + * + * Twin of Python ``mesh_dims_of_cells`` and of the long-range operator for a + * single frame: along every lattice vector ``a_d`` the mesh holds the + * smallest even number of points without a prime factor above 7 of at least + * ``ceil(|a_d| / spacing)``. The squared length sums the squares of the + * components, each rounded on its own as Python forms it; the explicit fma of + * a square with zero is that rounded square, which the compiler does not + * fuse into the sum, so that the three size the mesh identically whatever + * the floating-point contraction of the build. + * + * @param box The row-major cell, the rows its lattice vectors, nine values in + * the length unit of the model. + * @param spacing The largest mesh spacing, in the length unit of the model. + * @throw deepmd::deepmd_exception If the mesh exceeds 2^20 points along a + * lattice vector, the limit of the operator. + */ +inline std::array meshDims(const std::vector& box, + const double spacing) { + std::array dims = {}; + for (int d = 0; d < 3; ++d) { + const double* v = box.data() + 3 * d; + const double length = + std::sqrt(std::fma(v[0], v[0], 0.0) + std::fma(v[1], v[1], 0.0) + + std::fma(v[2], v[2], 0.0)); + const double points = std::ceil(length / spacing); + if (!(points < double(1 << 20))) { + throw deepmd::deepmd_exception( + "the mesh of the cell exceeds 2^20 points along a lattice vector; " + "a lattice vector of " + + std::to_string(length) + " at a mesh spacing of " + + std::to_string(spacing) + " needs " + std::to_string(points) + "."); + } + dims[d] = evenFftSize(static_cast(points)); + } + return dims; +} + /** * @brief Build the mesh topology a periodic long-range module of the graph * lower reads. * * Twin of Python ``build_mesh_grid`` in * ``deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py`` for a single - * frame: along every lattice vector ``a_d`` the mesh holds the smallest even - * number of points without a prime factor above 7 of at least - * ``ceil(|a_d| / spacing)``, with the length formed term by term as Python - * forms it, so that both size the mesh identically. Without a cell, the frame - * is non-periodic: the cell is the identity, the mesh minimal and the frame + * frame, on the mesh of ``meshDims``. Without a cell, the frame is + * non-periodic: the cell is the identity, the mesh minimal and the frame * masked. * * @param coord Node positions, shape ``(N, 3)``, on the model device. @@ -128,12 +174,7 @@ inline std::vector buildMeshTopology( for (int k = 0; k < 9; ++k) { cell[k] = box[k]; } - for (int d = 0; d < 3; ++d) { - const double* v = cell.data() + 3 * d; - const double length = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); - dims[d] = - evenFftSize(static_cast(std::ceil(length / spacing))); - } + dims = meshDims(box, spacing); } const double volume = std::abs(cell[0] * (cell[4] * cell[8] - cell[5] * cell[7]) - @@ -156,5 +197,100 @@ inline std::vector buildMeshTopology( return fields; } +/** + * @brief The cell input of the device-resident evaluations of a model with a + * long-range module, in pinned host memory. + * + * The compact lower reads the cell on the device while the stream runs its + * evaluation, so that a CUDA graph capturing the evaluation reads the cell + * of every replay from the same memory. One event orders those reads: an + * evaluation on another stream than the previous one first waits on the + * device for the previous one, so that the event of the last evaluation + * follows every earlier read, and a cell other than the one held waits for + * that event before the memory is written. + */ +class PinnedCell { + public: + PinnedCell() = default; + PinnedCell(const PinnedCell&) = delete; + PinnedCell& operator=(const PinnedCell&) = delete; + // The destructor reports no error. + ~PinnedCell() { +#ifdef DEEPMD_TORCH_HAS_CUDA + if (read_ != nullptr) { + int current = 0; + cudaGetDevice(¤t); + cudaSetDevice(device_); + cudaEventDestroy(read_); + cudaSetDevice(current); + } +#endif + } + + /** + * @brief Hold ``box`` for the next evaluation, which is enqueued on the + * current stream of ``device``. + * + * @param box The row-major cell, nine values, or empty for a non-periodic + * frame, whose cell is zero. + * @param device Index of the CUDA device of the evaluation. + * @return The cell input, a pinned float64 CPU tensor of shape ``(1, 9)``. + */ + const torch::Tensor& hold(const std::vector& box, const int device) { + if (!cell_.defined()) { + cell_ = torch::zeros( + {1, 9}, + torch::TensorOptions().dtype(torch::kFloat64).pinned_memory(true)); + } +#ifdef DEEPMD_TORCH_HAS_CUDA + device_ = static_cast(device); + stream_ = c10::cuda::getCurrentCUDAStream(device_).stream(); + if (read_ != nullptr && stream_ != read_stream_) { + C10_CUDA_CHECK(cudaStreamWaitEvent(stream_, read_, 0)); + } +#else + (void)device; +#endif + std::array values = {}; + std::copy(box.begin(), box.end(), values.begin()); + double* const cell = cell_.data_ptr(); + if (std::memcmp(cell, values.data(), sizeof(values)) != 0) { +#ifdef DEEPMD_TORCH_HAS_CUDA + if (read_ != nullptr) { + C10_CUDA_CHECK(cudaEventSynchronize(read_)); + } +#endif + std::memcpy(cell, values.data(), sizeof(values)); + } + return cell_; + } + + /** + * @brief Mark the evaluation of the cell last held, once enqueued, as the + * last that reads the cell. + */ + void record() { +#ifdef DEEPMD_TORCH_HAS_CUDA + const c10::cuda::CUDAGuard device_guard(device_); + if (read_ == nullptr) { + C10_CUDA_CHECK(cudaEventCreateWithFlags(&read_, cudaEventDisableTiming)); + } + C10_CUDA_CHECK(cudaEventRecord(read_, stream_)); + read_stream_ = stream_; +#endif + } + + private: + torch::Tensor cell_; +#ifdef DEEPMD_TORCH_HAS_CUDA + c10::DeviceIndex device_ = 0; + // The stream of the evaluation of the cell last held. + cudaStream_t stream_ = nullptr; + // The last evaluation that reads the cell, and its stream. + cudaEvent_t read_ = nullptr; + cudaStream_t read_stream_ = nullptr; +#endif +}; + } // namespace ptexpt } // namespace deepmd diff --git a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc index b45ea775b3..9658effbc1 100644 --- a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc @@ -27,6 +27,12 @@ #include #include "commonPT.h" +#ifdef DEEPMD_TORCH_HAS_CUDA +#include +#include +#include +#include +#endif #endif namespace { @@ -379,21 +385,27 @@ TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, deepmd::deepmd_exception); } -#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT && defined(DEEPMD_TORCH_HAS_CUDA) // The device-resident entry point of the compressed archives, which reads the -// positions of the nodes and the cell of a periodic system besides the -// compact canonical graph. A non-periodic system that one rank holds whole is -// evaluated eagerly, captured into a CUDA graph and replayed; a periodic one -// is evaluated eagerly, on the local atoms with the ghost atoms folded onto -// their owners, as one rank builds the graph, or with a node for every ghost -// atom, an image of its owner. Every evaluation writes into outputs reset to -// NaN, so an unwritten output fails the comparison. +// positions of the nodes and the cell of a periodic system besides the compact +// canonical graph. A graph without halo nodes is evaluated eagerly, captured +// into a CUDA graph and replayed: that of a non-periodic system that one rank +// holds whole, and that of a periodic system on the local atoms with the ghost +// atoms folded onto their owners, as one rank builds it, whose replay evaluates +// the cell of its call while the mesh of the cell keeps its size. A periodic +// graph with a node for every ghost atom, an image of its owner, is evaluated +// eagerly. Every evaluation writes into outputs reset to NaN, so an unwritten +// output fails the comparison, and a replay allocates no device memory, which +// tells it from an eager evaluation or a capture. class TestInferDeepPotDpa4cLongRangeDevicePtExpt : public ::testing::Test { protected: // A replay whose capture fails evaluates eagerly instead, with the same // results; only this warning tells the two apart. static constexpr const char* kCaptureWarning = "could not capture the evaluation"; + // Clock cycles for which at::cuda::sleep holds a stream back: tens of + // milliseconds, far longer than enqueueing an evaluation takes. + static constexpr std::int64_t kHoldCycles = std::int64_t{1} << 27; std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, @@ -419,13 +431,23 @@ class TestInferDeepPotDpa4cLongRangeDevicePtExpt : public ::testing::Test { }; // The outputs of an evaluation on the host, the forces folded onto the - // owners of the nodes. + // owners of the nodes, and the number of blocks that the caching allocator + // handed out on the device during the evaluation. struct Results { std::vector atom_energy; std::vector force; std::vector virial; + std::int64_t allocations = 0; }; + static std::int64_t device_allocations() { + const auto stats = c10::cuda::CUDACachingAllocator::getDeviceStats(0); + return stats + .allocation[static_cast( + c10::CachingAllocator::StatType::AGGREGATE)] + .allocated; + } + void load(const char* model_path, const char* ref_path) { if (!torch::cuda::is_available()) { GTEST_SKIP() << "Skip because the compact canonical path needs a CUDA " @@ -440,7 +462,9 @@ class TestInferDeepPotDpa4cLongRangeDevicePtExpt : public ::testing::Test { dp.init(model_path); ASSERT_TRUE(dp.uses_canonical_graph_inference()); ASSERT_TRUE(dp.has_long_range()); - ASSERT_TRUE(dp.canonical_graph_requires_source_csr()); + // The fused long-range operator serves the charge network of one channel + // and reads the source CSR for the per-atom virial alone. + ASSERT_FALSE(dp.canonical_graph_requires_source_csr()); ref.load(ref_path); testing::internal::CaptureStderr(); capturing_stderr = true; @@ -500,9 +524,36 @@ class TestInferDeepPotDpa4cLongRangeDevicePtExpt : public ::testing::Test { return graph; } - // Evaluates ``graph`` through compute_canonical_graph_gpu on the current - // stream, in the cell ``cell`` (null for a non-periodic system). - Results evaluate(Graph& graph, const int nloc, const double* cell) { + // ``graph`` with outputs of its own. + static Graph with_own_outputs(const Graph& graph) { + Graph copy = graph; + copy.atom_energy = torch::empty_like(graph.atom_energy); + copy.force = torch::empty_like(graph.force); + copy.virial = torch::empty_like(graph.virial); + copy.atom_virial = torch::empty_like(graph.atom_virial); + return copy; + } + + // The cell scaled by ``factor``. + std::vector scaled_box(const double factor) const { + std::vector cell = box; + for (double& entry : cell) { + entry *= factor; + } + return cell; + } + + // Enqueues on the current stream an evaluation of ``graph`` through + // compute_canonical_graph_gpu, in the cell ``cell`` (null for a + // non-periodic system), into outputs reset to NaN; ``stream`` is that + // stream, or null, with which the call waits for it before returning. + // Without ``source_csr`` the call passes neither the source CSR nor the + // per-atom virial. + void enqueue(Graph& graph, + const int nloc, + const double* cell, + void* stream, + const bool source_csr = true) { const double nan = std::numeric_limits::quiet_NaN(); graph.atom_energy.fill_(nan); graph.force.fill_(nan); @@ -511,33 +562,58 @@ class TestInferDeepPotDpa4cLongRangeDevicePtExpt : public ::testing::Test { const deepmd::CanonicalGraphTensorPack& compact = graph.tensors; dp.compute_canonical_graph_gpu( graph.atom_energy.data_ptr(), graph.force.data_ptr(), - graph.virial.data_ptr(), graph.atom_virial.data_ptr(), + graph.virial.data_ptr(), + source_csr ? graph.atom_virial.data_ptr() : nullptr, compact.atype.data_ptr(), compact.source.data_ptr(), deepmd::canonical_index_width(graph.nodes), compact.edge_vec.data_ptr(), compact.destination_row_ptr.data_ptr(), - compact.source_row_ptr.data_ptr(), - compact.source_order.data_ptr(), + source_csr ? compact.source_row_ptr.data_ptr() : nullptr, + source_csr ? compact.source_order.data_ptr() : nullptr, graph.coord.data_ptr(), /*world=*/nullptr, cell, nloc, - graph.nodes, compact.source.size(0)); + graph.nodes, compact.source.size(0), stream); + } + + // The atomic energies, forces and virial of an evaluation of ``graph`` on + // the host, the forces folded onto the owners of the nodes. + static Results collect(const Graph& graph, + const int nloc, + const torch::Tensor& atom_energy, + const torch::Tensor& force, + const torch::Tensor& virial) { const auto host_vector = [](const torch::Tensor& tensor) { const torch::Tensor values = tensor.cpu().contiguous(); return std::vector(values.data_ptr(), values.data_ptr() + values.numel()); }; Results results; - results.atom_energy = host_vector(graph.atom_energy); - results.virial = host_vector(graph.virial); - const std::vector force = host_vector(graph.force); + results.atom_energy = host_vector(atom_energy); + results.virial = host_vector(virial); + const std::vector node_force = host_vector(force); results.force.assign(3 * std::size_t(nloc), 0.0); for (int node = 0; node < graph.nodes; ++node) { for (int dd = 0; dd < 3; ++dd) { - results.force[3 * graph.owner[node] + dd] += force[3 * node + dd]; + results.force[3 * graph.owner[node] + dd] += node_force[3 * node + dd]; } } return results; } + // Evaluates ``graph`` on the current stream, in the cell ``cell`` (null for + // a non-periodic system), with or without the source CSR. + Results evaluate(Graph& graph, + const int nloc, + const double* cell, + const bool source_csr = true) { + const std::int64_t allocations = device_allocations(); + enqueue(graph, nloc, cell, /*stream=*/nullptr, source_csr); + const std::int64_t allocated = device_allocations() - allocations; + Results results = + collect(graph, nloc, graph.atom_energy, graph.force, graph.virial); + results.allocations = allocated; + return results; + } + // Agreement of an evaluation with the reference of a section. void expect_reference(const Results& results, const char* section) const { LongRangeRef expected; @@ -548,6 +624,42 @@ class TestInferDeepPotDpa4cLongRangeDevicePtExpt : public ::testing::Test { } expected.check(energy, results.force, results.virial, 1e-4); } + + // Evaluates the folded graph of the system, periodic in its cell or not, + // with the source CSR and then without it, which declines the per-atom + // virial alone: the energies, forces and virial agree bit for bit, and the + // evaluations without it are captured and replayed. + void expect_the_source_csr_optional(const bool periodic) { + const int nloc = static_cast(atype.size()); + std::vector> nlist; + std::vector coord_ext = coord; + std::vector atype_ext = atype; + std::vector mapping(nloc); + for (int ii = 0; ii < nloc; ++ii) { + mapping[ii] = ii; + } + if (periodic) { + _build_nlist(nlist, coord_ext, atype_ext, mapping, coord, atype, + box, static_cast(dp.cutoff())); + } else { + nlist = nopbc_nlist(coord, dp.cutoff()); + } + Graph graph = + build_graph(nlist, coord_ext, atype_ext, mapping, nloc, /*fold=*/true); + const double* cell = periodic ? box.data() : nullptr; + const Results with_csr = evaluate(graph, nloc, cell); + expect_reference(with_csr, periodic ? "periodic" : "default"); + for (int call = 0; call < 3; ++call) { + const Results without = evaluate(graph, nloc, cell, /*source_csr=*/false); + const std::string what = + "call " + std::to_string(call) + " without the source CSR"; + EXPECT_EQ(without.atom_energy, with_csr.atom_energy) << what; + EXPECT_EQ(without.force, with_csr.force) << what; + EXPECT_EQ(without.virial, with_csr.virial) << what; + EXPECT_EQ(without.allocations == 0, call == 2) + << what << ": " << without.allocations << " device allocations"; + } + } }; TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, cluster_is_replayed) { @@ -563,14 +675,18 @@ TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, cluster_is_replayed) { Graph graph = build_graph(nopbc_nlist(coord, dp.cutoff()), coord, atype, mapping, nloc, /*fold=*/true); ASSERT_EQ(graph.nodes, nloc); - // Eager, captured and replayed. + // Eager, captured, and replayed from the third call on. const Results eager = evaluate(graph, nloc, nullptr); expect_reference(eager, "default"); - for (int call = 1; call < 3; ++call) { - const Results replayed = evaluate(graph, nloc, nullptr); - EXPECT_EQ(replayed.atom_energy, eager.atom_energy) << "call " << call; - EXPECT_EQ(replayed.force, eager.force) << "call " << call; - EXPECT_EQ(replayed.virial, eager.virial) << "call " << call; + EXPECT_GT(eager.allocations, 0); + for (int call = 1; call < 4; ++call) { + const Results again = evaluate(graph, nloc, nullptr); + EXPECT_EQ(again.atom_energy, eager.atom_energy) << "call " << call; + EXPECT_EQ(again.force, eager.force) << "call " << call; + EXPECT_EQ(again.virial, eager.virial) << "call " << call; + EXPECT_EQ(again.allocations == 0, call > 1) + << "call " << call << ": " << again.allocations + << " device allocations"; } // The positions are part of the key of the replay: the same values at // another address are evaluated anew. @@ -578,6 +694,7 @@ TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, cluster_is_replayed) { const Results moved = evaluate(graph, nloc, nullptr); EXPECT_EQ(moved.atom_energy, eager.atom_energy); EXPECT_EQ(moved.force, eager.force); + EXPECT_GT(moved.allocations, 0); } TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, periodic_cell) { @@ -592,12 +709,152 @@ TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, periodic_cell) { _build_nlist(nlist, coord_ext, atype_ext, mapping, coord, atype, box, static_cast(dp.cutoff())); const int nloc = static_cast(atype.size()); - for (const bool fold : {true, false}) { - Graph graph = build_graph(nlist, coord_ext, atype_ext, mapping, nloc, fold); - ASSERT_EQ(graph.nodes, fold ? nloc : static_cast(atype_ext.size())); - for (int call = 0; call < 2; ++call) { - expect_reference(evaluate(graph, nloc, box.data()), "periodic"); + // At the mesh spacing of the model, the mesh of the cell has 40 points + // along every lattice vector, as it has for the cell strained by 1%, while + // the cell stretched by 10% has a mesh of 42, 42 and 48 points. + const std::vector strained = scaled_box(1.01); + const std::vector stretched = scaled_box(1.1); + const auto expect_equal = [](const Results& actual, const Results& expected, + const bool replayed, const std::string& what) { + EXPECT_EQ(actual.atom_energy, expected.atom_energy) << what; + EXPECT_EQ(actual.force, expected.force) << what; + EXPECT_EQ(actual.virial, expected.virial) << what; + EXPECT_EQ(actual.allocations == 0, replayed) + << what << ": " << actual.allocations << " device allocations"; + }; + + // The folded graph: the strained cell is evaluated eagerly, and the first + // cell, under the same key, is captured and replayed. The strained cell is + // then replayed at once, since its mesh keeps the size of the capture. + Graph graph = build_graph(nlist, coord_ext, atype_ext, mapping, nloc, true); + ASSERT_EQ(graph.nodes, nloc); + const Results strained_eager = evaluate(graph, nloc, strained.data()); + EXPECT_GT(strained_eager.allocations, 0); + const Results first = evaluate(graph, nloc, box.data()); + expect_reference(first, "periodic"); + EXPECT_GT(first.allocations, 0); + EXPECT_NE(first.atom_energy, strained_eager.atom_energy); + expect_equal(evaluate(graph, nloc, box.data()), first, true, "replay"); + expect_equal(evaluate(graph, nloc, strained.data()), strained_eager, true, + "replay in the strained cell"); + expect_equal(evaluate(graph, nloc, box.data()), first, true, + "replay back in the first cell"); + + // A mesh of another size is evaluated eagerly until it has held for 64 + // calls, then once more under its new key, captured at the 66th call and + // replayed from the 67th on. + const Results stretched_eager = evaluate(graph, nloc, stretched.data()); + EXPECT_GT(stretched_eager.allocations, 0); + EXPECT_NE(stretched_eager.atom_energy, first.atom_energy); + for (int call = 2; call <= 67; ++call) { + expect_equal(evaluate(graph, nloc, stretched.data()), stretched_eager, + call == 67, + "call " + std::to_string(call) + " in the stretched cell"); + } + + // The unfolded graph has halo nodes and is evaluated eagerly. + Graph unfolded = + build_graph(nlist, coord_ext, atype_ext, mapping, nloc, false); + ASSERT_EQ(unfolded.nodes, static_cast(atype_ext.size())); + const Results unfolded_first = evaluate(unfolded, nloc, box.data()); + expect_reference(unfolded_first, "periodic"); + for (int call = 1; call < 3; ++call) { + expect_equal(evaluate(unfolded, nloc, box.data()), unfolded_first, false, + "call " + std::to_string(call) + " of the unfolded graph"); + } +} + +TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, + a_new_cell_waits_for_the_evaluations) { + ASSERT_NO_FATAL_FAILURE( + load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath)); + if (IsSkipped()) { + return; + } + std::vector> nlist; + std::vector coord_ext; + std::vector atype_ext, mapping; + _build_nlist(nlist, coord_ext, atype_ext, mapping, coord, atype, box, + static_cast(dp.cutoff())); + const int nloc = static_cast(atype.size()); + // The strained cell keeps the mesh of the cell. + const std::vector strained = scaled_box(1.01); + // Three graphs over the same inputs with outputs of their own: every call + // changes the key of the replay, so that every evaluation runs eagerly and + // no capture or release of a graph waits for the device. + Graph held_graph = + build_graph(nlist, coord_ext, atype_ext, mapping, nloc, /*fold=*/true); + Graph same_graph = with_own_outputs(held_graph); + Graph strained_graph = with_own_outputs(held_graph); + const c10::cuda::CUDAStream held = + c10::cuda::getStreamFromPool(/*isHighPriority=*/false, /*device=*/0); + const c10::cuda::CUDAStream other = + c10::cuda::getStreamFromPool(/*isHighPriority=*/false, /*device=*/0); + ASSERT_NE(held.stream(), other.stream()); + // The cell on a held stream, the same cell on another stream, then the + // strained cell on that stream, each evaluation followed by copies of its + // outputs on its stream. + const auto run = [&](const bool hold) { + std::vector copies; + const auto run_on = [&](Graph& graph, const std::vector& cell, + const c10::cuda::CUDAStream& stream) { + const c10::cuda::CUDAStreamGuard guard(stream); + enqueue(graph, nloc, cell.data(), stream.stream()); + for (const torch::Tensor& output : + {graph.atom_energy, graph.force, graph.virial}) { + copies.push_back(output.clone()); + } + }; + if (hold) { + const c10::cuda::CUDAStreamGuard guard(held); + at::cuda::sleep(kHoldCycles); } + run_on(held_graph, box, held); + run_on(same_graph, box, other); + run_on(strained_graph, strained, other); + torch::cuda::synchronize(); + return copies; + }; + // The first round runs the one-time work of the program. In the second, + // the held stream runs its evaluation after the strained cell has been + // handed over, and that call writes the cell only once the evaluations + // before it, the held one included, have read it. + run(/*hold=*/false); + const std::vector copies = run(/*hold=*/true); + const Results first = evaluate(held_graph, nloc, box.data()); + expect_reference(first, "periodic"); + const Results first_strained = evaluate(held_graph, nloc, strained.data()); + const auto expect_equal = [&](const Graph& graph, const std::size_t at, + const Results& expected, + const std::string& what) { + const Results actual = + collect(graph, nloc, copies[at], copies[at + 1], copies[at + 2]); + EXPECT_EQ(actual.atom_energy, expected.atom_energy) << what; + EXPECT_EQ(actual.force, expected.force) << what; + EXPECT_EQ(actual.virial, expected.virial) << what; + }; + expect_equal(held_graph, 0, first, "the cell on the held stream"); + expect_equal(same_graph, 3, first, "the cell on the other stream"); + expect_equal(strained_graph, 6, first_strained, + "the strained cell on the other stream"); +} + +TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, + the_source_csr_is_optional_in_a_cluster) { + ASSERT_NO_FATAL_FAILURE(load(kCompressedModelPath, kCompressedRefPath)); + if (IsSkipped()) { + return; + } + expect_the_source_csr_optional(/*periodic=*/false); +} + +TEST_F(TestInferDeepPotDpa4cLongRangeDevicePtExpt, + the_source_csr_is_optional_in_a_cell) { + ASSERT_NO_FATAL_FAILURE( + load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath)); + if (IsSkipped()) { + return; } + expect_the_source_csr_optional(/*periodic=*/true); } #endif diff --git a/source/op/pt/dpa4c/graph_compress.cu b/source/op/pt/dpa4c/graph_compress.cu index 8fc7b53e01..8eef579e48 100644 --- a/source/op/pt/dpa4c/graph_compress.cu +++ b/source/op/pt/dpa4c/graph_compress.cu @@ -26,6 +26,7 @@ #include "../canonical_index.h" #include "../graph_ops.h" +#include "../long_range/charge_sum.h" #include "graph_compress_launch.h" namespace { @@ -1090,11 +1091,12 @@ __device__ __forceinline__ double fixed_point_value(unsigned long long word) { // Per-node preparation of one evaluation: the node types with virtual atoms // (negative types) mapped onto type zero, the ownership mask that selects the // energy-contributing nodes -- a real atom among the leading ``n_local`` -// nodes of its frame -- as the int32 output mask and as the fp64 cotangent -// that seeds the fitting backward, and cleared fixed-point force totals. The -// frame of a node comes from ``frame_row_ptr`` when the batch holds several -// frames and is frame zero otherwise; nodes past the last frame are padding, -// which no frame owns. +// nodes of its frame -- as the int32 output mask, as the fp64 cotangent that +// seeds the fitting backward and, unless ``charged`` is null, as the boolean +// mask of the nodes that carry the charges of a long-range module, and +// cleared fixed-point force totals. The frame of a node comes from +// ``frame_row_ptr`` when the batch holds several frames and is frame zero +// otherwise; nodes past the last frame are padding, which no frame owns. __global__ void energy_force_prologue_kernel( long node_count, long frame_count, @@ -1104,6 +1106,7 @@ __global__ void energy_force_prologue_kernel( long* __restrict__ node_type, double* __restrict__ seed, int* __restrict__ mask, + bool* __restrict__ charged, unsigned long long* __restrict__ force, unsigned long long* __restrict__ spin_force, const float* __restrict__ contact_radius, @@ -1142,6 +1145,9 @@ __global__ void energy_force_prologue_kernel( const bool owned = raw >= 0 && node - frame_begin < n_local[frame]; seed[node] = owned ? 1.0 : 0.0; mask[node] = owned ? 1 : 0; + if (charged != nullptr) { + charged[node] = owned; + } #pragma unroll for (int word = 0; word < kForceWords; ++word) { force[node * kForceWords + word] = 0ull; @@ -1174,6 +1180,11 @@ constexpr int kEpilogueTotals = 10; // double-precision work, which most devices run at a small fraction of the // single-precision rate, spreads over kEpilogueTotals times as many warps. // The warps of a total also write the per-node results that pair with it. +// +// A long-range module adds its node energies, its forces at fixed charges and +// its per-node virials, ``extra_energy``, ``extra_force`` and +// ``extra_virial``, onto those of the descriptor and the fitting; all three +// are null for a model without one. __global__ void __launch_bounds__(kEpilogueThreads* kEpilogueTotals) energy_force_epilogue_kernel( long node_count, @@ -1184,6 +1195,9 @@ __global__ void __launch_bounds__(kEpilogueThreads* kEpilogueTotals) const double* __restrict__ raw_energy, const double* __restrict__ seed, const float* __restrict__ node_virial, + const double* __restrict__ extra_energy, + const float* __restrict__ extra_force, + const float* __restrict__ extra_virial, double* __restrict__ force, double* __restrict__ force_mag, double* __restrict__ atom_energy, @@ -1209,10 +1223,12 @@ __global__ void __launch_bounds__(kEpilogueThreads* kEpilogueTotals) // A node whose force contributions left the range of the fixed-point // totals reports its energy as not a number, which marks the forces of // its frame. - const double value = - force_fixed[node * kForceWords + kForceRangeWord] != 0 - ? CUDART_NAN - : raw_energy[node] * seed[node]; + double value = force_fixed[node * kForceWords + kForceRangeWord] != 0 + ? CUDART_NAN + : raw_energy[node] * seed[node]; + if (extra_energy != nullptr) { + value += extra_energy[node]; + } atom_energy[node] = value; if (node < frame_end) { sum += value; @@ -1220,8 +1236,12 @@ __global__ void __launch_bounds__(kEpilogueThreads* kEpilogueTotals) continue; } if (component < 3) { - force[node * 3 + component] = + double value = fixed_point_value(force_fixed[node * kForceWords + component]); + if (extra_force != nullptr) { + value += static_cast(extra_force[node * 3 + component]); + } + force[node * 3 + component] = value; if (force_mag != nullptr) { // The magnetic force is minus the energy gradient with respect to the // moment: its on-site part and the neighbour part of its sources. @@ -1235,6 +1255,9 @@ __global__ void __launch_bounds__(kEpilogueThreads* kEpilogueTotals) } if (node < frame_end) { sum += static_cast(node_virial[node * 9 + component]); + if (extra_virial != nullptr) { + sum += static_cast(extra_virial[node * 9 + component]); + } } } @@ -1313,7 +1336,9 @@ __global__ void __launch_bounds__(kEpilogueThreads* kEpilogueTotals) } // Per-node results of a batch of several frames; the frame totals follow -// from the segment reductions of the caller. +// from the segment reductions of the caller. ``extra_energy`` and +// ``extra_force`` are the contributions of a long-range module, as in +// energy_force_epilogue_kernel. __global__ void energy_force_nodes_kernel( long node_count, const unsigned long long* __restrict__ force_fixed, @@ -1321,6 +1346,8 @@ __global__ void energy_force_nodes_kernel( const float* __restrict__ spin_gradient, const double* __restrict__ raw_energy, const double* __restrict__ seed, + const double* __restrict__ extra_energy, + const float* __restrict__ extra_force, double* __restrict__ force, double* __restrict__ force_mag, double* __restrict__ atom_energy, @@ -1329,17 +1356,25 @@ __global__ void energy_force_nodes_kernel( node < node_count; node += static_cast(blockDim.x) * gridDim.x) { #pragma unroll for (int component = 0; component < 3; ++component) { - force[node * 3 + component] = + double value = fixed_point_value(force_fixed[node * kForceWords + component]); + if (extra_force != nullptr) { + value += static_cast(extra_force[node * 3 + component]); + } + force[node * 3 + component] = value; if (force_mag != nullptr) { force_mag[node * 3 + component] = -(static_cast(spin_gradient[node * 3 + component]) + fixed_point_value(spin_fixed[node * kForceWords + component])); } } - atom_energy[node] = force_fixed[node * kForceWords + kForceRangeWord] != 0 - ? CUDART_NAN - : raw_energy[node] * seed[node]; + double value = force_fixed[node * kForceWords + kForceRangeWord] != 0 + ? CUDART_NAN + : raw_energy[node] * seed[node]; + if (extra_energy != nullptr) { + value += extra_energy[node]; + } + atom_energy[node] = value; if (atom_virial_zero != nullptr) { #pragma unroll for (int component = 0; component < 9; ++component) { @@ -1355,6 +1390,112 @@ int epilogue_blocks(long node_count) { kEpilogueBlocks))); } +// Validates the graph inputs of a fused energy-force operator named ``op`` +// that its prologue reads before the runs validate the rest of the payload: +// the frame counts, the node types and, for a windowed descriptor, the +// contact radii, and the source CSR, which is either complete, requesting the +// per-atom virial, or the two-entry placeholder that declines it. ``source`` +// is replaced by its typed view. Returns whether the per-atom virial is +// requested. +bool check_energy_force_inputs(const char* op, + const torch::Tensor& edge_vec, + torch::Tensor& source, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_row_ptr, + const torch::Tensor& source_order, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& atype, + const torch::Tensor& type_embedding, + const torch::Tensor& contact_radius, + double f_inner, + double f_outer) { + TORCH_CHECK(edge_vec.is_cuda(), op, ": edge_vec must be a CUDA tensor"); + source = deepmd::canonical_index_view(source); + TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), op, + ": source and edge_vec must share the edge axis"); + const long node_count = atype.size(0); + const long edge_count = edge_vec.size(0); + TORCH_CHECK(destination_row_ptr.numel() == node_count + 1, op, + ": atype and destination_row_ptr describe different node " + "counts"); + const torch::Device device = edge_vec.device(); + const auto on_device = [&](const torch::Tensor& tensor) { + return tensor.is_cuda() && tensor.device() == device && + tensor.is_contiguous(); + }; + TORCH_CHECK(on_device(n_node) && n_node.dim() == 1 && + n_node.scalar_type() == torch::kInt64 && on_device(n_local) && + n_local.dim() == 1 && + n_local.scalar_type() == torch::kInt64 && + n_local.numel() == n_node.numel() && n_node.numel() > 0, + op, + ": n_node and n_local must be contiguous int64 tensors on the " + "device of edge_vec with one shared shape (nf,), nf >= 1"); + TORCH_CHECK(on_device(atype) && atype.dim() == 1 && + atype.scalar_type() == torch::kInt64, + op, + ": atype must be a contiguous int64 tensor of shape (N,) on the " + "device of edge_vec"); + TORCH_CHECK( + f_outer <= f_inner || (on_device(contact_radius) && + contact_radius.scalar_type() == torch::kFloat32 && + contact_radius.dim() == 1 && + contact_radius.size(0) == type_embedding.size(0)), + op, + ": a windowed descriptor needs contiguous fp32 contact radii of shape " + "(type_count,) on the device of edge_vec"); + const bool atom_virial = source_row_ptr.numel() == node_count + 1 && + source_order.numel() == edge_count; + TORCH_CHECK( + atom_virial || (source_row_ptr.numel() == 2 && source_order.numel() == 2), + op, ": the source CSR must hold ", node_count + 1, " row pointers and ", + edge_count, + " positions, or two entries each to decline the per-atom virial; got ", + source_row_ptr.numel(), " and ", source_order.numel()); + if (atom_virial) { + TORCH_CHECK((source.scalar_type() == torch::kUInt32 || + source.scalar_type() == torch::kInt64) && + on_device(source_row_ptr) && + source_row_ptr.scalar_type() == torch::kInt64 && + on_device(source_order) && + source_order.scalar_type() == torch::kInt64, + op, + ": the per-atom virial needs uint32 or int64 sources and a " + "source CSR of contiguous int64 row pointers and positions on " + "the device of edge_vec"); + } + return atom_virial; +} + +// A fixed-point force word cannot hold a non-finite value, so a frame of a +// batch whose energy or virial is not finite reports every force of its +// nodes, and every magnetic force unless ``force_mag`` is null, as not a +// number (see energy_force_epilogue_kernel). Nodes past the last frame +// belong to none and keep their forces. +void mark_broken_frames(const torch::Tensor& energy, + const torch::Tensor& virial, + const torch::Tensor& frame_row_ptr, + long node_count, + torch::Tensor& force, + torch::Tensor* force_mag) { + const auto broken = + torch::cat({torch::logical_not(torch::isfinite(energy).all(1).logical_and( + torch::isfinite(virial).flatten(1).all(1))), + torch::zeros({1}, energy.options().dtype(torch::kBool))}); + const auto frame = + torch::searchsorted(frame_row_ptr, + torch::arange(node_count, frame_row_ptr.options()), + /*out_int32=*/false, /*right=*/true) - + 1; + const auto node_broken = broken.index_select(0, frame).unsqueeze(1); + force.masked_fill_(node_broken, std::numeric_limits::quiet_NaN()); + if (force_mag != nullptr) { + force_mag->masked_fill_(node_broken, + std::numeric_limits::quiet_NaN()); + } +} + } // namespace // Energy, force and virial of one compressed inference step, evaluated over @@ -1451,71 +1592,12 @@ dpa4c_canonical_compress_energy_force( int64_t tile, torch::Tensor pair_table, std::optional> fitting_routes) { - TORCH_CHECK(edge_vec.is_cuda(), - "dpa4c_canonical_compress_energy_force: edge_vec must be a CUDA " - "tensor"); - source = deepmd::canonical_index_view(source); - TORCH_CHECK(source.dim() == 1 && source.numel() == edge_vec.size(0), - "dpa4c_canonical_compress_energy_force: source and edge_vec " - "must share the edge axis"); + const bool atom_virial = check_energy_force_inputs( + "dpa4c_canonical_compress_energy_force", edge_vec, source, + destination_row_ptr, source_row_ptr, source_order, n_node, n_local, atype, + type_embedding, contact_radius, f_inner, f_outer); const c10::cuda::CUDAGuard device_guard(edge_vec.device()); const long node_count = atype.size(0); - const long edge_count = edge_vec.size(0); - TORCH_CHECK(destination_row_ptr.numel() == node_count + 1, - "dpa4c_canonical_compress_energy_force: atype and " - "destination_row_ptr describe different node counts"); - // The prologue reads the frame counts, the node types and, for a windowed - // descriptor, the contact radii before the runs validate the rest of the - // payload, so those inputs are validated here. - const torch::Device device = edge_vec.device(); - const auto on_device = [&](const torch::Tensor& tensor) { - return tensor.is_cuda() && tensor.device() == device && - tensor.is_contiguous(); - }; - TORCH_CHECK(on_device(n_node) && n_node.dim() == 1 && - n_node.scalar_type() == torch::kInt64 && on_device(n_local) && - n_local.dim() == 1 && - n_local.scalar_type() == torch::kInt64 && - n_local.numel() == n_node.numel() && n_node.numel() > 0, - "dpa4c_canonical_compress_energy_force: n_node and n_local must " - "be contiguous int64 tensors on the device of edge_vec with one " - "shared shape (nf,), nf >= 1"); - TORCH_CHECK(on_device(atype) && atype.dim() == 1 && - atype.scalar_type() == torch::kInt64, - "dpa4c_canonical_compress_energy_force: atype must be a " - "contiguous int64 tensor of shape (N,) on the device of " - "edge_vec"); - TORCH_CHECK( - f_outer <= f_inner || (on_device(contact_radius) && - contact_radius.scalar_type() == torch::kFloat32 && - contact_radius.dim() == 1 && - contact_radius.size(0) == type_embedding.size(0)), - "dpa4c_canonical_compress_energy_force: a windowed descriptor " - "needs contiguous fp32 contact radii of shape (type_count,) on " - "the device of edge_vec"); - // The source CSR is either complete, which requests the per-atom virial, - // or the two-entry placeholder, which declines it. - const bool atom_virial = source_row_ptr.numel() == node_count + 1 && - source_order.numel() == edge_count; - TORCH_CHECK( - atom_virial || (source_row_ptr.numel() == 2 && source_order.numel() == 2), - "dpa4c_canonical_compress_energy_force: the source CSR must " - "hold ", - node_count + 1, " row pointers and ", edge_count, - " positions, or two entries each to decline the per-atom " - "virial; got ", - source_row_ptr.numel(), " and ", source_order.numel()); - if (atom_virial) { - TORCH_CHECK((source.scalar_type() == torch::kUInt32 || - source.scalar_type() == torch::kInt64) && - on_device(source_row_ptr) && - source_row_ptr.scalar_type() == torch::kInt64 && - on_device(source_order) && - source_order.scalar_type() == torch::kInt64, - "dpa4c_canonical_compress_energy_force: the per-atom virial " - "needs uint32 or int64 sources and a source CSR of contiguous " - "int64 row pointers and positions on the device of edge_vec"); - } const int channels = static_cast(type_embedding.size(1)); const bool has_spin = spin.dim() == 2; const Dimensions widths = @@ -1571,7 +1653,7 @@ dpa4c_canonical_compress_energy_force( node_count, frame_count, n_local.data_ptr(), frame_count > 1 ? frame_row_ptr.data_ptr() : nullptr, atype.data_ptr(), node_type.data_ptr(), - seed.data_ptr(), mask.data_ptr(), + seed.data_ptr(), mask.data_ptr(), nullptr, reinterpret_cast(force_fixed.data_ptr()), has_spin ? reinterpret_cast( spin_fixed.data_ptr()) @@ -1731,7 +1813,8 @@ dpa4c_canonical_compress_energy_force( node_count, n_node.data_ptr(), force_words, spin_words, has_spin ? spin_gradient.data_ptr() : nullptr, raw_energy.data_ptr(), seed.data_ptr(), - node_virial.data_ptr(), force.data_ptr(), + node_virial.data_ptr(), nullptr, nullptr, nullptr, + force.data_ptr(), has_spin ? force_mag.data_ptr() : nullptr, atom_energy.data_ptr(), atom_virial_zero, partial.data_ptr(), evaluation_words, energy.data_ptr(), @@ -1741,32 +1824,15 @@ dpa4c_canonical_compress_energy_force( energy_force_nodes_kernel<<>>( node_count, force_words, spin_words, has_spin ? spin_gradient.data_ptr() : nullptr, - raw_energy.data_ptr(), seed.data_ptr(), - force.data_ptr(), + raw_energy.data_ptr(), seed.data_ptr(), nullptr, + nullptr, force.data_ptr(), has_spin ? force_mag.data_ptr() : nullptr, atom_energy.data_ptr(), atom_virial_zero); DPA4C_CHECK_LAUNCH("dpa4c energy-force nodes"); energy.copy_(frame_scalar_sum(atom_energy, n_node)); virial.copy_(frame_virial_sum(node_virial, n_node)); - // A fixed-point force word cannot hold a non-finite value, so a frame - // whose energy or virial is not finite reports every force as not a - // number (see energy_force_epilogue_kernel). - // Nodes past the last frame belong to none and keep their forces. - const auto broken = torch::cat( - {torch::logical_not(torch::isfinite(energy).all(1).logical_and( - torch::isfinite(virial).flatten(1).all(1))), - torch::zeros({1}, energy.options().dtype(torch::kBool))}); - const auto frame = - torch::searchsorted(frame_row_ptr, torch::arange(node_count, i64), - /*out_int32=*/false, /*right=*/true) - - 1; - const auto node_broken = broken.index_select(0, frame); - force.masked_fill_(node_broken.unsqueeze(1), - std::numeric_limits::quiet_NaN()); - if (has_spin) { - force_mag.masked_fill_(node_broken.unsqueeze(1), - std::numeric_limits::quiet_NaN()); - } + mark_broken_frames(energy, virial, frame_row_ptr, node_count, force, + has_spin ? &force_mag : nullptr); } if (atom_virial) { atom_virial_output.copy_( @@ -1777,6 +1843,388 @@ dpa4c_canonical_compress_energy_force( force_mag, mask}; } +// Energy, force and virial of one compressed inference step of a model with a +// latent-charge long-range module. +// +// The module reads charges that its network predicts from the descriptor of +// every node, and the energy of the charges depends on all of them, so the +// descriptor backward cannot start before every charge is known. The +// operator evaluates the descriptor of every node with its moment state, the +// charges from the descriptor, and the long range of the charges at fixed +// charges with its derivative with respect to them +// (deepmd::latent_charge_sum); then the energy fitting with its cotangent and +// the cotangent of the charge network, seeded with that derivative, which it +// adds; and one descriptor backward from the sum of the two cotangents. The +// two networks evaluate the node axis in runs of ``tile`` nodes, as in +// dpa4c_canonical_compress_energy_force, which bounds the memory of their +// activations by the run, the pre-activations of the charge network, which +// its backward reads, by the node axis. The edge backward adds every force +// onto the fixed-point totals as in dpa4c_canonical_compress_energy_force, +// and the epilogue adds the energies, the forces at fixed charges and the +// per-node virials of the long range. +// +// The charge network has one channel and the layout of an eligible energy +// fitting; its per-type bias is ``charge_bias``. ``coord``, ``total_charge``, +// ``box``, ``coefficient``, ``communicator``, ``kernel``, ``channel``, +// ``prefactor``, ``order`` and ``spacing`` are those of +// deepmd::latent_charge_sum, whose charged nodes are the owned real nodes of +// the frames. The per-atom virial and the source CSR follow +// dpa4c_canonical_compress_energy_force; the descriptor carries no spin +// condition and the model no analytical pair potential. +std::tuple +dpa4c_canonical_compress_long_range_energy_force( + torch::Tensor edge_vec, + torch::Tensor source, + torch::Tensor destination_row_ptr, + torch::Tensor source_row_ptr, + torch::Tensor source_order, + torch::Tensor n_node, + torch::Tensor n_local, + torch::Tensor atype, + torch::Tensor table, + torch::Tensor basis_table, + torch::Tensor basis_matrix, + torch::Tensor pair_film, + torch::Tensor pair_mixing, + torch::Tensor type_embedding, + torch::Tensor readout_matrices, + torch::Tensor coupling_meta, + torch::Tensor coupling_entry, + torch::Tensor coupling_value, + torch::Tensor output_mean, + torch::Tensor output_inv_std, + torch::Tensor spin, + torch::Tensor spin_pair, + torch::Tensor spin_type, + torch::Tensor contact_radius, + int64_t lmax, + double table_stride, + double table_max, + double rcut, + double eps, + double degree_floor, + double f_inner, + double f_outer, + std::vector ws, + std::vector bs, + std::vector resnets, + torch::Tensor w_head, + torch::Tensor b_head, + torch::Tensor bias_atom_e, + int64_t act, + int64_t tile, + std::vector charge_ws, + std::vector charge_bs, + std::vector charge_resnets, + torch::Tensor charge_w_head, + torch::Tensor charge_b_head, + torch::Tensor charge_bias, + int64_t charge_act, + torch::Tensor coord, + torch::Tensor total_charge, + std::optional box, + torch::Tensor coefficient, + torch::Tensor communicator, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing, + std::optional> fitting_routes) { + constexpr const char* kOp = + "dpa4c_canonical_compress_long_range_energy_force"; + const bool atom_virial = check_energy_force_inputs( + kOp, edge_vec, source, destination_row_ptr, source_row_ptr, source_order, + n_node, n_local, atype, type_embedding, contact_radius, f_inner, f_outer); + const c10::cuda::CUDAGuard device_guard(edge_vec.device()); + TORCH_CHECK(spin.dim() != 2, kOp, + ": a descriptor with a spin condition is not supported"); + const long node_count = atype.size(0); + TORCH_CHECK(coord.is_cuda() && coord.device() == edge_vec.device() && + coord.dim() == 2 && coord.size(0) == node_count && + coord.size(1) == 3 && coord.scalar_type() == torch::kFloat64, + kOp, + ": coord must be a float64 tensor of shape (N, 3) on the device " + "of edge_vec"); + const int channels = static_cast(type_embedding.size(1)); + const Dimensions widths = profile_dimensions(channels, static_cast(lmax), + /*has_spin=*/false); + const long frame_count = n_node.size(0); + const auto f32 = edge_vec.options().dtype(torch::kFloat32); + const auto f64 = edge_vec.options().dtype(torch::kFloat64); + const auto i64 = edge_vec.options().dtype(torch::kInt64); + auto edge_vec_float = edge_vec.to(torch::kFloat32).contiguous(); + const auto stream = at::cuda::getCurrentCUDAStream(); + + // Every output receives its own allocation, so that no two of the six + // unannotated results share storage. + auto energy = torch::empty({frame_count, 1}, f64); + auto atom_energy = torch::empty({node_count, 1}, f64); + auto force = torch::empty({node_count, 3}, f64); + auto virial = torch::empty({frame_count, 3, 3}, f64); + auto atom_virial_output = torch::empty({node_count, 3, 3}, f32); + auto mask = + torch::empty({node_count}, edge_vec.options().dtype(torch::kInt32)); + if (node_count == 0) { + energy.zero_(); + virial.zero_(); + return {energy, atom_energy, force, virial, atom_virial_output, mask}; + } + + auto raw_energy = torch::empty({node_count}, f64); + auto seed = torch::empty({node_count}, f64); + auto charged = + torch::empty({node_count}, edge_vec.options().dtype(torch::kBool)); + auto node_type = torch::empty({node_count}, i64); + auto force_fixed = torch::empty({node_count, kForceWords}, i64); + auto node_virial = torch::empty({node_count, 9}, f32); + auto edge_gradient = + atom_virial ? torch::empty_like(edge_vec_float) : torch::Tensor(); + const int blocks = epilogue_blocks(node_count); + auto words = + torch::zeros({static_cast(sizeof(EvaluationWords) / sizeof(int))}, + edge_vec.options().dtype(torch::kInt32)); + auto* evaluation_words = + reinterpret_cast(words.data_ptr()); + const torch::Tensor frame_row_ptr = + frame_count > 1 ? frame_row_pointer(n_node) : torch::Tensor(); + energy_force_prologue_kernel<<>>( + node_count, frame_count, n_local.data_ptr(), + frame_count > 1 ? frame_row_ptr.data_ptr() : nullptr, + atype.data_ptr(), node_type.data_ptr(), + seed.data_ptr(), mask.data_ptr(), charged.data_ptr(), + reinterpret_cast(force_fixed.data_ptr()), + nullptr, f_outer > f_inner ? contact_radius.data_ptr() : nullptr, + evaluation_words); + DPA4C_CHECK_LAUNCH("dpa4c long-range energy-force prologue"); + + const FittingLayerPlan plan = fitting_layer_plan(ws); + const FittingLayerPlan charge_plan = fitting_layer_plan(charge_ws); + TORCH_CHECK(plan.n_layer > 0 && ws[0].size(0) == widths.output_width && + charge_plan.n_layer > 0 && + charge_ws[0].size(0) == widths.output_width, + kOp, + ": the first weights of the energy fitting and of the charge " + "network must match the descriptor width"); + TORCH_CHECK(charge_w_head.numel() == charge_ws.back().size(1) && + charge_bias.dim() == 1 && + charge_bias.size(0) == bias_atom_e.size(0), + kOp, + ": the charge network must have a scalar head and one bias per " + "type, as the energy fitting has"); + + // === Step 1. Descriptor and moment state of every node === + auto descriptor = torch::empty({node_count, widths.output_width}, f32); + auto state = torch::empty({node_count, widths.moment_width + 2}, f32); + auto empty_index = torch::empty({0}, source.options()); + auto empty_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); + const Payload payload{source, + empty_mask, + empty_index, + destination_row_ptr, + node_type, + table, + basis_table, + basis_matrix, + pair_film, + pair_mixing, + type_embedding, + readout_matrices, + coupling_meta, + coupling_entry, + coupling_value, + output_mean, + output_inv_std, + spin, + spin_pair, + spin_type, + contact_radius, + true, + lmax, + table_stride, + table_max, + rcut, + eps, + degree_floor, + f_inner, + f_outer, + 0}; + Arguments arguments = + build_arguments(payload, edge_vec_float, channels, widths); + arguments.system_node_count = node_count; + arguments.descriptor = descriptor.data_ptr(); + arguments.state_out = state.data_ptr(); + arguments.type_mismatch = &evaluation_words->node_types.type_mismatch; + arguments.largest_radius = &evaluation_words->node_types.largest_radius; + arguments.type_present = + present_types(arguments, &evaluation_words->node_types); + dispatch(channels, false, arguments, stream); + + // === Step 2. Charges of the nodes, run by run === + // The charge network reads the descriptor before the energy fitting + // overwrites it with its cotangent. The pre-activations of every run stay, + // run after run in the layout of their run, for the backward of the run. + const long run = tile > 0 + ? std::max(1, std::min(tile, node_count)) + : node_count; + const long width = widths.output_width; + const long charge_saved_width = charge_plan.saved_width(); + auto charge_saved = torch::empty({node_count * charge_saved_width}, f32); + const int charge_slots = charge_plan.n_layer > 1 ? 2 : 1; + auto charge_scratch = + torch::empty({charge_slots, run, charge_plan.width_max}, f32); + float* const charge_slot[2] = {charge_scratch[0].data_ptr(), + charge_slots > 1 + ? charge_scratch[1].data_ptr() + : charge_scratch[0].data_ptr()}; + auto raw_charge = torch::empty({node_count}, f64); + const auto charge_bias_double = charge_bias.to(torch::kFloat64).contiguous(); + for (long begin = 0; begin < node_count; begin += run) { + const long count = std::min(run, node_count - begin); + fitting_forward_range( + stream, charge_plan, descriptor.data_ptr() + begin * width, + width, node_type.data_ptr() + begin, charge_ws, charge_bs, + charge_resnets, charge_w_head, charge_b_head, charge_bias_double, + charge_act, count, + charge_saved.data_ptr() + begin * charge_saved_width, + charge_slot, raw_charge.data_ptr() + begin); + } + + // === Step 3. Long range of the charges === + // The per-node virials of the charged nodes sum to the virial of each frame, + // which the frame totals form with those of the descriptor. + const auto long_range = deepmd::long_range::latent_charge_sum( + coord, raw_charge.to(torch::kFloat32).unsqueeze(1), charged, n_node, + total_charge, box, coefficient, communicator, kernel, channel, prefactor, + order, spacing); + const torch::Tensor& long_range_energy = std::get<0>(long_range); + const torch::Tensor& charge_gradient = std::get<1>(long_range); + const torch::Tensor& long_range_force = std::get<2>(long_range); + const torch::Tensor& long_range_atom_virial = std::get<3>(long_range); + + // === Step 4. Cotangents of the two networks, run by run === + // The full runs and the tail run each take the fitting route of their own + // size, as in dpa4c_canonical_compress_energy_force. + const std::vector routes = + std::move(fitting_routes).value_or(std::vector{}); + const long tail = node_count % run; + const long full_count = node_count / run; + const int64_t full_route = fitting_route(plan, ws, run, routes); + const int64_t tail_route = + tail > 0 ? fitting_route(plan, ws, tail, routes) : full_route; + const long per_layer_nodes = + std::max(full_count > 0 && full_route == kFittingPerLayer ? run : 0L, + tail > 0 && tail_route == kFittingPerLayer ? tail : 0L); + auto saved = torch::empty({run * plan.saved_width()}, f32); + torch::Tensor transposed; + torch::Tensor scratch; + float* slot[2] = {nullptr, nullptr}; + if ((full_count > 0 && full_route != kFittingPerLayer) || + (tail > 0 && tail_route != kFittingPerLayer)) { + transposed = torch::empty({fitting_fused_transposed_floats(ws)}, f32); + fitting_fused_transpose(stream, plan, ws, transposed.data_ptr()); + } + if (per_layer_nodes > 0) { + const int slots = plan.n_layer > 1 ? 2 : 1; + scratch = torch::empty({slots, per_layer_nodes, plan.width_max}, f32); + slot[0] = scratch[0].data_ptr(); + slot[1] = slots > 1 ? scratch[1].data_ptr() : slot[0]; + } + const auto charge_seed = + charge_gradient.select(1, 0).to(torch::kFloat64).contiguous(); + auto charge_cotangent = torch::empty({run, width}, f32); + for (long begin = 0; begin < node_count; begin += run) { + const long count = std::min(run, node_count - begin); + // The cotangent of the energy fitting replaces the descriptor of the run, + // which the charge network has read. + float* const run_descriptor = descriptor.data_ptr() + begin * width; + const int64_t route = count == run ? full_route : tail_route; + if (route != kFittingPerLayer) { + fitting_energy_gradient_range( + stream, plan, run_descriptor, width, + node_type.data_ptr() + begin, ws, bs, resnets, w_head, b_head, + bias_atom_e, act, transposed.data_ptr(), + seed.data_ptr() + begin, count, saved.data_ptr(), + raw_energy.data_ptr() + begin, run_descriptor, route); + } else { + fitting_forward_range(stream, plan, run_descriptor, width, + node_type.data_ptr() + begin, ws, bs, resnets, + w_head, b_head, bias_atom_e, act, count, + saved.data_ptr(), slot, + raw_energy.data_ptr() + begin); + fitting_backward_range( + stream, plan, seed.data_ptr() + begin, + saved.data_ptr(), ws, bs, resnets, w_head, act, count, slot[0], + plan.n_layer > 1 ? slot[1] : nullptr, run_descriptor); + } + fitting_backward_range( + stream, charge_plan, charge_seed.data_ptr() + begin, + charge_saved.data_ptr() + begin * charge_saved_width, charge_ws, + charge_bs, charge_resnets, charge_w_head, charge_act, count, + charge_slot[0], charge_plan.n_layer > 1 ? charge_slot[1] : nullptr, + charge_cotangent.data_ptr()); + descriptor.slice(0, begin, begin + count) + .add_(charge_cotangent.slice(0, 0, count)); + } + + // === Step 5. Descriptor backward and force assembly === + arguments.descriptor_gradient = descriptor.data_ptr(); + arguments.state = state.data_ptr(); + arguments.moment_gradient = state.data_ptr(); + arguments.force = + reinterpret_cast(force_fixed.data_ptr()); + arguments.node_virial = node_virial.data_ptr(); + arguments.edge_gradient = + atom_virial ? edge_gradient.data_ptr() : nullptr; + arguments.clear_padding = true; + dispatch(channels, true, arguments, stream); + + // === Step 6. Node results and frame totals === + const auto* force_words = reinterpret_cast( + force_fixed.data_ptr()); + float* atom_virial_zero = + atom_virial ? nullptr : atom_virial_output.data_ptr(); + if (frame_count == 1) { + auto partial = torch::empty({blocks, kEpilogueTotals}, f64); + energy_force_epilogue_kernel<<>>( + node_count, n_node.data_ptr(), force_words, nullptr, nullptr, + raw_energy.data_ptr(), seed.data_ptr(), + node_virial.data_ptr(), long_range_energy.data_ptr(), + long_range_force.data_ptr(), + long_range_atom_virial.data_ptr(), force.data_ptr(), + nullptr, atom_energy.data_ptr(), atom_virial_zero, + partial.data_ptr(), evaluation_words, energy.data_ptr(), + virial.data_ptr()); + DPA4C_CHECK_LAUNCH("dpa4c long-range energy-force epilogue"); + } else { + energy_force_nodes_kernel<<>>( + node_count, force_words, nullptr, nullptr, + raw_energy.data_ptr(), seed.data_ptr(), + long_range_energy.data_ptr(), + long_range_force.data_ptr(), force.data_ptr(), nullptr, + atom_energy.data_ptr(), atom_virial_zero); + DPA4C_CHECK_LAUNCH("dpa4c long-range energy-force nodes"); + energy.copy_(frame_scalar_sum(atom_energy, n_node)); + node_virial.add_(long_range_atom_virial); + virial.copy_(frame_virial_sum(node_virial, n_node)); + mark_broken_frames(energy, virial, frame_row_ptr, node_count, force, + nullptr); + } + if (atom_virial) { + atom_virial_output.copy_( + canonical_source_virial(edge_gradient, edge_vec_float, source_row_ptr, + source_order, node_count)); + atom_virial_output.add_(long_range_atom_virial.reshape({node_count, 3, 3})); + } + return {energy, atom_energy, force, virial, atom_virial_output, mask}; +} + TORCH_LIBRARY_FRAGMENT(deepmd, library) { library.impl("dpa4c_graph_compress", torch::kCUDA, &dpa4c_graph_compress); library.impl("dpa4c_graph_compress_backward", torch::kCUDA, @@ -1789,4 +2237,6 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { &dpa4c_canonical_compress_backward_inplace); library.impl("dpa4c_canonical_compress_energy_force", torch::kCUDA, &dpa4c_canonical_compress_energy_force); + library.impl("dpa4c_canonical_compress_long_range_energy_force", torch::kCUDA, + &dpa4c_canonical_compress_long_range_energy_force); } diff --git a/source/op/pt/dpa4c/ops.cc b/source/op/pt/dpa4c/ops.cc index 4ea1c73334..2bd4a0c04e 100644 --- a/source/op/pt/dpa4c/ops.cc +++ b/source/op/pt/dpa4c/ops.cc @@ -34,6 +34,17 @@ // CSR; a two-entry placeholder in place of `source_row_ptr` and // `source_order` declines it, which leaves it zero. // +// The fused long-range operator evaluates a model whose atomic energy adds a +// latent-charge long-range module with a charge network of one channel. It +// takes the arguments of the fused operator up to `tile`, the network that +// predicts the charges from the descriptor (`charge_ws` to `charge_act`, in +// the layout of the energy fitting, with the per-type bias `charge_bias`), +// the inputs of deepmd::latent_charge_sum (`coord` to `spacing`, see +// source/op/pt/long_range/ops.cc), whose charged nodes are the owned real +// nodes, and the fitting routes. It returns the outputs of the fused operator +// but the magnetic force: the descriptor carries no spin condition and the +// model no pair potential. +// // Every schema carries, next to the radial table, the low-rank table and the // expansion that a single-type scan reads instead of the scalar-only channels // (empty together when the descriptor admits none; see `low_rank_basis`). @@ -126,4 +137,25 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "int tile, Tensor pair_table, int[]? fitting_routes=None) " "-> (Tensor energy, Tensor atom_energy, Tensor force, Tensor virial, " "Tensor atom_virial, Tensor force_mag, Tensor mask)"); + library.def( + "dpa4c_canonical_compress_long_range_energy_force(Tensor edge_vec, " + "Tensor source, Tensor destination_row_ptr, Tensor source_row_ptr, " + "Tensor source_order, Tensor n_node, Tensor n_local, Tensor atype, " + "Tensor table, Tensor basis_table, Tensor basis_matrix, " + "Tensor pair_film, Tensor pair_mixing, Tensor type_embedding, " + "Tensor readout_matrices, Tensor coupling_meta, Tensor coupling_entry, " + "Tensor coupling_value, Tensor output_mean, Tensor output_inv_std, " + "Tensor spin, Tensor spin_pair, Tensor spin_type, Tensor contact_radius, " + "int lmax, float table_stride, float table_max, float rcut, float eps, " + "float degree_floor, float f_inner, float f_outer, " + "Tensor[] ws, Tensor[] bs, int[] resnets, " + "Tensor w_head, Tensor b_head, Tensor bias_atom_e, int act, int tile, " + "Tensor[] charge_ws, Tensor[] charge_bs, int[] charge_resnets, " + "Tensor charge_w_head, Tensor charge_b_head, Tensor charge_bias, " + "int charge_act, Tensor coord, Tensor total_charge, Tensor? box, " + "Tensor coefficient, Tensor communicator, int kernel, int channel, " + "float prefactor, int order, float spacing, " + "int[]? fitting_routes=None) " + "-> (Tensor energy, Tensor atom_energy, Tensor force, Tensor virial, " + "Tensor atom_virial, Tensor mask)"); } diff --git a/source/op/pt/long_range/charge_sum.h b/source/op/pt/long_range/charge_sum.h index 04245d3287..4d439a250c 100644 --- a/source/op/pt/long_range/charge_sum.h +++ b/source/op/pt/long_range/charge_sum.h @@ -64,5 +64,22 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, int owner_count, const torch::Tensor* communicator); +// deepmd::latent_charge_sum (schema and formulas in ops.cc): validates the +// inputs and evaluates the frames by the direct or the mesh sum, across the +// domains of ``communicator`` when it holds several ranks. +ChargeSumOutputs latent_charge_sum(torch::Tensor coord, + torch::Tensor charge, + torch::Tensor node_mask, + torch::Tensor n_node_per_frame, + torch::Tensor total_charge, + c10::optional box, + torch::Tensor coefficient, + torch::Tensor communicator, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing); + } // namespace long_range } // namespace deepmd diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu index 2e0f128d85..cb7c56ae93 100644 --- a/source/op/pt/long_range/direct_sum.cu +++ b/source/op/pt/long_range/direct_sum.cu @@ -25,7 +25,12 @@ // A thread sums one tile in float32 and adds the tile sums with compensated // float32 arithmetic. The partner axis is split into contiguous chunks, so // that a launch fills the device several times over; every chunk writes its -// sums, and a second pass adds the chunks in a fixed order in float64. +// sums, and a second pass adds the chunks in a fixed order in float64. A +// system too small to fill the device with tiles of kWideTile nodes takes +// tiles of one warp, so that its pairs spread over more multiprocessors. Its +// chunks hold kNarrowChunkTiles tiles each, counted from the first partner of +// each block, rather than equal parts of the partner range: the chunks of a +// block then do not depend on the partners past the frames of its nodes. // // The charge constraint and the per-frame virial reduce over the nodes of // every frame, in slices whose float64 sums are added in a fixed order, so @@ -54,15 +59,29 @@ namespace { -// Nodes per block, which is also the length of a tile of partners. +// Threads per block of the per-node kernels. constexpr int kThreads = 128; +// Nodes of a tile of the pair kernel: a block owns one tile of nodes and +// walks the partners of its chunk in tiles of the same length. Every thread +// sums the pairs of a tile in float32 and adds the tile sum to its +// compensated sums. +constexpr int kWideTile = 128; +constexpr int kNarrowTile = 32; // Partners evaluated together, which share every coefficient load. constexpr int kBatch = 4; constexpr int kMaxGaussians = 32; // Fills of the resident block slots of the device that one launch aims for. constexpr int kWaves = 8; -// Shortest chunk of the partner axis, in tiles. -constexpr int kMinChunkTiles = 2; +// Tiles per chunk of wide tiles where a frame that spans every partner takes +// the most chunks it may: pair_launch rounds the chunk count up from its tile +// count over this. +constexpr int kWideChunkTiles = 2; +// Tiles of a chunk of narrow tiles: two, so that a chunk that holds a tile +// entirely past the partners of a frame holds at most one tile with partners +// of that frame. The compensated sum of one tile carries no error, so that +// the zero sums of the tiles past the frame leave it unchanged, as the frame +// on its own would. +constexpr int kNarrowChunkTiles = 2; enum KernelKind : int { kLes = 0, kSog = 1 }; @@ -315,7 +334,7 @@ struct OwnNode { // Adds the unscaled potential, force direction and virial that the partners // of one staged tile give the node to ``sum``. -template +template __device__ __forceinline__ void accumulate_tile( const float4* __restrict__ tile_head, const float4* __restrict__ tile_tail, @@ -326,7 +345,7 @@ __device__ __forceinline__ void accumulate_tile( const OwnNode& own, float (&sum)[sum_width()]) { #pragma unroll 4 - for (int base = 0; base < kThreads; base += kBatch) { + for (int base = 0; base < kTile; base += kBatch) { float dx[kBatch]; float dy[kBatch]; float dz[kBatch]; @@ -403,9 +422,12 @@ __device__ __forceinline__ void accumulate_tile( // Sums, for every owner node of one block and every partner of one chunk, the // unscaled potential, force direction and virial of the node, and writes them // to ``partial[chunk][component][node - owner_begin]``. The owners are the -// nodes [owner_begin, owner_begin + owner_count). -template -__global__ void __launch_bounds__(kThreads) +// nodes [owner_begin, owner_begin + owner_count). A chunk is the part +// blockIdx.y of gridDim.y equal parts of the partner range of the block or, +// with ``chunk_tiles`` > 0, its chunk_tiles tiles from tile blockIdx.y * +// chunk_tiles of the range. +template +__global__ void __launch_bounds__(kTile) direct_sum_kernel(const float4* __restrict__ head, const float4* __restrict__ tail, const float* __restrict__ charge, @@ -415,22 +437,23 @@ __global__ void __launch_bounds__(kThreads) int owner_count, const double* __restrict__ coefficient, int gaussian_count, + int chunk_tiles, float* __restrict__ partial) { constexpr int kWidth = sum_width(); - __shared__ float4 tile_head[kThreads]; - __shared__ float4 tile_tail[kThreads]; - __shared__ float tile_charge[kThreads * kChannels]; + __shared__ float4 tile_head[kTile]; + __shared__ float4 tile_tail[kTile]; + __shared__ float tile_charge[kTile * kChannels]; __shared__ float4 gaussians[kMaxGaussians]; __shared__ int range[2]; __shared__ bool single_frame; const int owner_end = owner_begin + owner_count; - const int first = owner_begin + blockIdx.x * kThreads; + const int first = owner_begin + blockIdx.x * kTile; const int node = first + threadIdx.x; const int own_frame = node < owner_end ? frame[node] : -1; if constexpr (kKernel == kSog) { - for (int k = threadIdx.x; k < gaussian_count; k += kThreads) { + for (int k = threadIdx.x; k < gaussian_count; k += kTile) { const float amplitude = static_cast(coefficient[k]); const float rate = static_cast(coefficient[gaussian_count + k]); gaussians[k] = make_float4(amplitude, -rate, @@ -441,7 +464,7 @@ __global__ void __launch_bounds__(kThreads) // to: a contiguous range, since frames are contiguous on the node axis and // padding nodes trail them. if (threadIdx.x == 0) { - int last = min(first + kThreads, owner_end) - 1; + int last = min(first + kTile, owner_end) - 1; while (last > first && frame[last] < 0) { --last; } @@ -452,19 +475,20 @@ __global__ void __launch_bounds__(kThreads) // A block whose nodes all lie in one frame finds every partner of its // range in that frame, which spares the frame test of every pair. single_frame = low_frame >= 0 && low_frame == high_frame && - min(first + kThreads, owner_end) - 1 == last; + min(first + kTile, owner_end) - 1 == last; } __syncthreads(); // The chunk of the partner range this block sums, in whole tiles. - const int64_t tile_count = (range[1] - range[0] + kThreads - 1) / kThreads; - const int begin = - range[0] + - static_cast(tile_count * blockIdx.y / gridDim.y) * kThreads; - const int end = min( - range[1], - range[0] + static_cast(tile_count * (blockIdx.y + 1) / gridDim.y) * - kThreads); + const int64_t tile_count = (range[1] - range[0] + kTile - 1) / kTile; + int64_t first_tile = tile_count * blockIdx.y / gridDim.y; + int64_t end_tile = tile_count * (blockIdx.y + 1) / gridDim.y; + if (chunk_tiles > 0) { + first_tile = min(tile_count, int64_t{blockIdx.y} * chunk_tiles); + end_tile = min(tile_count, first_tile + chunk_tiles); + } + const int begin = range[0] + static_cast(first_tile) * kTile; + const int end = min(range[1], range[0] + static_cast(end_tile) * kTile); OwnNode own; own.index = node; @@ -484,7 +508,7 @@ __global__ void __launch_bounds__(kThreads) } CompensatedSum total[kWidth]; - for (int start = begin; start < end; start += kThreads) { + for (int start = begin; start < end; start += kTile) { const int partner = start + threadIdx.x; const bool staged = partner < end; // A slot past the chunk carries a zero charge and contributes nothing. @@ -508,15 +532,15 @@ __global__ void __launch_bounds__(kThreads) // nodes hold a node itself, so the tests are confined to them and to the // blocks that span frames; the choice is uniform over the block. if (!single_frame) { - accumulate_tile( + accumulate_tile( tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, own, sum); - } else if (start < first + kThreads && first < start + kThreads) { - accumulate_tile( + } else if (start < first + kTile && first < start + kTile) { + accumulate_tile( tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, own, sum); } else { - accumulate_tile( + accumulate_tile( tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, own, sum); } @@ -676,26 +700,6 @@ __global__ void project_charges(const float* __restrict__ potential, atom_energy[node] = 0.5 * energy; } -// Chunks of the partner axis for the given counts of owners and partners. The -// blocks of one launch all carry the same work, so the launch fills the -// resident block slots of the device several times over, which bounds the -// idle share of the last wave; no chunk is shorter than kMinChunkTiles tiles -// of a frame that spans every partner. -int chunk_count_for(int owner_count, - int partner_count, - int resident_blocks_per_multiprocessor) { - const int multiprocessors = - at::cuda::getCurrentDeviceProperties()->multiProcessorCount; - const int owner_blocks = (owner_count + kThreads - 1) / kThreads; - const int partner_tiles = (partner_count + kThreads - 1) / kThreads; - const int64_t wanted_blocks = static_cast(kWaves) * - resident_blocks_per_multiprocessor * - multiprocessors; - const int wanted = - static_cast((wanted_blocks + owner_blocks - 1) / owner_blocks); - return std::max(1, std::min(wanted, partner_tiles / kMinChunkTiles)); -} - // The staged partners, one row per node, and the sums of the owners. struct DirectSumBuffers { const float4* head; @@ -711,6 +715,52 @@ struct DirectSumBuffers { float* atom_virial; }; +// The launch of the pair kernel: whether it takes narrow tiles, and its +// chunks of the partner axis, ``chunk_count`` of them and, for narrow tiles, +// ``chunk_tiles`` tiles each (direct_sum_kernel). +struct PairLaunch { + bool narrow; + int chunk_count; + int chunk_tiles; +}; + +// The launch for the given counts of owners and partners. The blocks of one +// launch of wide tiles all carry the same work, so the launch fills the +// resident block slots of the device several times over, which bounds the +// idle share of the last wave. A frame that spans every partner has at most +// one chunk per kWideChunkTiles of its tiles, rounded up, so that in a launch +// with slots to spare no chunk holds more than kWideChunkTiles tiles: the +// longest block sets the duration of such a launch. A system whose wide +// blocks would fill less than a third of the resident block slots leaves too +// few warps on the multiprocessors to hide the latency of the pair loop, and +// takes narrow tiles in chunks of kNarrowChunkTiles, as many as a frame that +// spans every partner holds. +PairLaunch pair_launch(int owner_count, + int partner_count, + int resident_blocks_per_multiprocessor) { + const int multiprocessors = + at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + const int owner_blocks = (owner_count + kWideTile - 1) / kWideTile; + const int partner_tiles = (partner_count + kWideTile - 1) / kWideTile; + const int64_t wanted_blocks = static_cast(kWaves) * + resident_blocks_per_multiprocessor * + multiprocessors; + const int wanted = + static_cast((wanted_blocks + owner_blocks - 1) / owner_blocks); + const int chunk_limit = + (partner_tiles + kWideChunkTiles - 1) / kWideChunkTiles; + const int chunk_count = std::max(1, std::min(wanted, chunk_limit)); + if (3 * int64_t{owner_blocks} * chunk_count >= + int64_t{resident_blocks_per_multiprocessor} * multiprocessors) { + return {false, chunk_count, 0}; + } + const int narrow_tiles = (partner_count + kNarrowTile - 1) / kNarrowTile; + return { + true, + std::max(1, (narrow_tiles + kNarrowChunkTiles - 1) / kNarrowChunkTiles), + kNarrowChunkTiles}; +} + template int resident_blocks_per_multiprocessor() { // The occupancy is a property of the compiled kernel and of the device; a @@ -718,42 +768,59 @@ int resident_blocks_per_multiprocessor() { static const int blocks = [] { int value = 0; C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &value, direct_sum_kernel, kThreads, 0)); + &value, direct_sum_kernel, kWideTile, + 0)); return std::max(1, value); }(); return blocks; } -template -void launch(const DirectSumBuffers& buffers, - int owner_begin, - int owner_count, - double prefactor, - cudaStream_t stream) { +template +void launch_tiles(const DirectSumBuffers& buffers, + int owner_begin, + int owner_count, + const PairLaunch& pair, + double prefactor, + cudaStream_t stream) { constexpr int kWidth = sum_width(); - const int chunk_count = - chunk_count_for(owner_count, buffers.node_count, - resident_blocks_per_multiprocessor()); - auto partial = - torch::empty({static_cast(chunk_count) * owner_count * kWidth}, - torch::TensorOptions() - .dtype(torch::kFloat32) - .device(torch::kCUDA, c10::cuda::current_device())); - const dim3 grid((owner_count + kThreads - 1) / kThreads, chunk_count); - direct_sum_kernel<<>>( + auto partial = torch::empty( + {static_cast(pair.chunk_count) * owner_count * kWidth}, + torch::TensorOptions() + .dtype(torch::kFloat32) + .device(torch::kCUDA, c10::cuda::current_device())); + const dim3 grid((owner_count + kTile - 1) / kTile, pair.chunk_count); + direct_sum_kernel<<>>( buffers.head, buffers.tail, buffers.charge, buffers.frame, buffers.frame_ptr, owner_begin, owner_count, buffers.coefficient, - buffers.gaussian_count, partial.data_ptr()); + buffers.gaussian_count, pair.chunk_tiles, partial.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); const int64_t reduce_threads = static_cast(kWidth) * owner_count; reduce_chunks <<((reduce_threads + 255) / 256), 256, 0, stream>>>( - partial.data_ptr(), owner_count, chunk_count, prefactor, + partial.data_ptr(), owner_count, pair.chunk_count, prefactor, kKernel == kLes ? buffers.coefficient : nullptr, buffers.potential, buffers.force, buffers.atom_virial); C10_CUDA_KERNEL_LAUNCH_CHECK(); } +template +void launch(const DirectSumBuffers& buffers, + int owner_begin, + int owner_count, + double prefactor, + cudaStream_t stream) { + const PairLaunch pair = + pair_launch(owner_count, buffers.node_count, + resident_blocks_per_multiprocessor()); + if (pair.narrow) { + launch_tiles( + buffers, owner_begin, owner_count, pair, prefactor, stream); + } else { + launch_tiles( + buffers, owner_begin, owner_count, pair, prefactor, stream); + } +} + template void dispatch_channels(int channels, const DirectSumBuffers& buffers, diff --git a/source/op/pt/long_range/latent_charge_sum.cu b/source/op/pt/long_range/latent_charge_sum.cu index 4763c72fcc..91af0c09b5 100644 --- a/source/op/pt/long_range/latent_charge_sum.cu +++ b/source/op/pt/long_range/latent_charge_sum.cu @@ -137,6 +137,11 @@ bool periodic_frames(const c10::optional& box) { return periodic > 0; } +} // namespace + +namespace deepmd { +namespace long_range { + ChargeSumOutputs latent_charge_sum(torch::Tensor coord, torch::Tensor charge, torch::Tensor node_mask, @@ -228,8 +233,10 @@ ChargeSumOutputs latent_charge_sum(torch::Tensor coord, n_node_per_frame, total_charge, communicator); } -} // namespace +} // namespace long_range +} // namespace deepmd TORCH_LIBRARY_FRAGMENT(deepmd, library) { - library.impl("latent_charge_sum", torch::kCUDA, &latent_charge_sum); + library.impl("latent_charge_sum", torch::kCUDA, + &deepmd::long_range::latent_charge_sum); } diff --git a/source/op/pt/long_range/mesh_sum.cu b/source/op/pt/long_range/mesh_sum.cu index 130cbdcb11..8bccbcbf90 100644 --- a/source/op/pt/long_range/mesh_sum.cu +++ b/source/op/pt/long_range/mesh_sum.cu @@ -57,6 +57,7 @@ #include #include +#include #include #include #include @@ -85,11 +86,20 @@ constexpr float kFixedScaleFloat = 1099511627776.f; // 2^40 constexpr double kPi = 3.14159265358979323846; -// The cell of one frame and its mesh. +// The cell of one frame and its mesh. The host sizes the mesh, which fixes +// the buffers and the launches of a call, and the device derives the +// geometry from the cell and the prefactor of the call, so that a CUDA graph +// capturing a call evaluates the cell of every replay on the meshes of the +// capture. struct FrameMesh { // B^-1, row-major: s_d = sum_j r_j inverse[3 j + d]. double inverse[9]; - double volume; + // 2 pi B^-1 in float32, row-major like inverse: the wave vector of the mode + // m is k_a = sum_d reciprocal[3 a + d] m_d. + float reciprocal[9]; + // The factor prefactor / V of the influence function in float32, V the + // volume of the cell. + float influence_scale; // -n_d (B^-1)_jd, row-major in (j, d): the force at fixed charges from the // derivatives of the energy with respect to the scaled coordinates u_d. float gradient_to_force[9]; @@ -116,7 +126,10 @@ int64_t even_fft_size(int64_t n) { } } -FrameMesh frame_mesh(const double* cell, double spacing) { +// The mesh of the cell ``cell`` at a spacing of at most ``spacing``, sized +// along every lattice vector; store_frame_mesh adds the geometry of the cell +// on the device. +FrameMesh frame_mesh_size(const double* cell, double spacing) { FrameMesh mesh{}; const double a = cell[0], b = cell[1], c = cell[2]; const double d = cell[3], e = cell[4], f = cell[5]; @@ -130,19 +143,16 @@ FrameMesh frame_mesh(const double* cell, double spacing) { std::sqrt(g * g + h * h + i * i); TORCH_CHECK(std::abs(det) > 1e-6 * lengths, "latent_charge_sum: a periodic cell is singular"); - const double adjugate[9] = {e * i - f * h, c * h - b * i, b * f - c * e, - f * g - d * i, a * i - c * g, c * d - a * f, - d * h - e * g, b * g - a * h, a * e - b * d}; - for (int k = 0; k < 9; ++k) { - mesh.inverse[k] = adjugate[k] / det; - } - mesh.volume = std::abs(det); - // The squared lengths of the lattice vectors are formed term by term, as - // the Python front end forms them, so that both round the mesh size - // identically. + // The squared length of a lattice vector sums the squares of its + // components, each rounded on its own as the Python front end and the C++ + // interface form them: the explicit fma of a square with zero is that + // rounded square, which the compiler does not fuse into the sum, so that + // all three size the mesh identically. for (int axis = 0; axis < 3; ++axis) { const double* v = cell + 3 * axis; - const double length = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + const double length = + std::sqrt(std::fma(v[0], v[0], 0.0) + std::fma(v[1], v[1], 0.0) + + std::fma(v[2], v[2], 0.0)); const double points = std::ceil(length / spacing); TORCH_CHECK(points < 1 << 20, "latent_charge_sum: the mesh of a cell exceeds 2^20 points " @@ -153,12 +163,6 @@ FrameMesh frame_mesh(const double* cell, double spacing) { TORCH_CHECK(int64_t{mesh.n[0]} * mesh.n[1] * mesh.n[2] < std::numeric_limits::max(), "latent_charge_sum: the mesh of a cell exceeds 2^31 points"); - for (int j = 0; j < 3; ++j) { - for (int d = 0; d < 3; ++d) { - mesh.gradient_to_force[3 * j + d] = - static_cast(-mesh.n[d] * mesh.inverse[3 * j + d]); - } - } return mesh; } @@ -396,6 +400,41 @@ __global__ void __launch_bounds__(kStencilWarps* kWarp) } } +// Stores the mesh of one frame, sized on the host and carried by value, with +// the geometry of its cell ``cell``. Thread k = 3 j + d of kGeometryThreads +// forms the entry (B^-1)_jd = C_dj / det B of the inverse of the cell B, C +// its cofactors, and the wave-vector and force factors of that entry; the +// first thread also forms the factor prefactor / V of the influence +// function. +constexpr int kGeometryThreads = 9; + +__global__ void store_frame_mesh(FrameMesh mesh, + const double* __restrict__ cell, + double prefactor, + FrameMesh* __restrict__ out) { + const int k = threadIdx.x; + const int j = k / 3; + const int d = k % 3; + const double det = cell[0] * (cell[4] * cell[8] - cell[5] * cell[7]) - + cell[1] * (cell[3] * cell[8] - cell[5] * cell[6]) + + cell[2] * (cell[3] * cell[7] - cell[4] * cell[6]); + const double* row1 = cell + 3 * ((d + 1) % 3); + const double* row2 = cell + 3 * ((d + 2) % 3); + const double inverse = (row1[(j + 1) % 3] * row2[(j + 2) % 3] - + row1[(j + 2) % 3] * row2[(j + 1) % 3]) / + det; + out->inverse[k] = inverse; + out->reciprocal[k] = static_cast(2.0 * kPi * inverse); + out->gradient_to_force[k] = static_cast(-mesh.n[d] * inverse); + if (k == 0) { + out->influence_scale = static_cast(prefactor / fabs(det)); + for (int axis = 0; axis < 3; ++axis) { + out->n[axis] = mesh.n[axis]; + } + out->offset = mesh.offset; + } +} + // Converts the fixed-point mesh to the float32 mesh that the transforms read. __global__ void fixed_to_float(int64_t count, const unsigned long long* __restrict__ fixed, @@ -414,7 +453,9 @@ struct IntegerWeights { }; // The inverse squared assignment factor |D_d(m)|^-2 along every axis of the -// half spectrum of one frame: n1 entries, then n2, then n3 / 2 + 1. +// half spectrum of a mesh: n1 entries, then n2, then n3 / 2 + 1. It depends +// only on the shape of the mesh and the order of the B-splines, so the plans +// of a shape hold it (create_mesh_plans). __global__ void assignment_factor(FrameMesh mesh, int order, IntegerWeights integer_weight, @@ -512,30 +553,30 @@ __device__ __forceinline__ void kernel_transform(int kernel, constexpr int kModeSums = 7; __global__ void __launch_bounds__(kThreads) - apply_influence(FrameMesh mesh, + apply_influence(const FrameMesh* __restrict__ frame, int channels, int kernel, const double* __restrict__ coefficient, int gaussians, - double prefactor, const float* __restrict__ factor, float2* __restrict__ spectrum, double* __restrict__ block_sum) { __shared__ KernelTerms terms; load_kernel_terms(kernel, coefficient, gaussians, terms); __syncthreads(); + const FrameMesh& mesh = *frame; const int n1 = mesh.n[0]; const int n2 = mesh.n[1]; const int n3 = mesh.n[2]; const int half = n3 / 2 + 1; const int modes = n1 * n2 * half; - // k_a = 2 pi sum_d (B^-1)_ad m_d + // k_a = sum_d reciprocal[3 a + d] m_d = 2 pi sum_d (B^-1)_ad m_d float reciprocal[9]; #pragma unroll for (int k = 0; k < 9; ++k) { - reciprocal[k] = static_cast(2.0 * kPi * mesh.inverse[k]); + reciprocal[k] = mesh.reciprocal[k]; } - const float scale = static_cast(prefactor / mesh.volume); + const float scale = mesh.influence_scale; // Every warp takes rows (i1, i2) of the half spectrum, its lanes the modes // of a row, which lie contiguously. const int lane = threadIdx.x % kWarp; @@ -944,29 +985,70 @@ IntegerWeights integer_bspline_weights(int order) { return weight; } +// The launch of a kernel of float64 work, one item per thread. The float64 +// pipe of a multiprocessor serves every warp on it, so few items spread over +// as many multiprocessors as there are warps of them, in blocks of a few +// warps, and many items fill blocks of kThreads. +struct Float64Launch { + int blocks; + int threads; +}; + +Float64Launch float64_launch(int64_t work) { + const int64_t multiprocessors = + at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + const int64_t warps = (work + kWarp - 1) / kWarp; + const int64_t warps_per_block = std::clamp( + (warps + multiprocessors - 1) / multiprocessors, 1, kThreads / kWarp); + const int threads = static_cast(warps_per_block * kWarp); + return {static_cast(std::max( + 1, std::min((work + threads - 1) / threads, 65535))), + threads}; +} + +// Blocks of the influence function of a mesh: a few per multiprocessor fill +// the device, and no more blocks than the rows of the half spectrum fill, a +// row to a warp, so that the block sums stay few. The count depends on the +// mesh and the device alone, so that the modes of a frame are summed in the +// same order whatever the other frames of its call. +int influence_blocks(const FrameMesh& mesh) { + constexpr int64_t kWarps = kThreads / kWarp; + const int64_t rows = int64_t{mesh.n[0]} * mesh.n[1]; + return static_cast(std::clamp( + (rows + kWarps - 1) / kWarps, 1, + 4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount)); +} + void check_cufft(cufftResult result, const char* operation) { TORCH_CHECK(result == CUFFT_SUCCESS, "latent_charge_sum: cuFFT ", operation, " failed with error ", static_cast(result)); } // The real-to-complex and complex-to-real transforms of the meshes of one -// shape, batched over the channels. The plans allocate no work area: every -// transform runs on a work area of the workspace of its call, set with the -// stream under ``mutex`` just before the transform is enqueued, so that -// concurrent calls never share one. +// shape, batched over the channels, and the inverse squared assignment +// factors of the shape at one B-spline order, which depend on nothing else +// and are formed once. The plans allocate no work area: every transform runs +// on a work area of the workspace of its call, set with the stream under +// ``mutex`` just before the transform is enqueued, so that concurrent calls +// never share one. A CUDA graph that captures a transform runs it on these +// plans at every replay, so it holds them while it lives: ``graphs`` counts +// the graphs that hold them. struct MeshPlans { cufftHandle forward = 0; cufftHandle inverse = 0; size_t work_bytes = 0; + // float32, n1 + n2 + n3 / 2 + 1 entries on the device (assignment_factor). + torch::Tensor factor; c10::DeviceIndex device = 0; std::mutex mutex; + std::atomic graphs{0}; MeshPlans() = default; MeshPlans(const MeshPlans&) = delete; MeshPlans& operator=(const MeshPlans&) = delete; - // A plan leaves the cache only past kPlanCacheSize shapes; its transforms - // may still be in flight, so the device finishes them first. The - // destructor reports no error. + // Plans leave the cache only past kPlanCacheSize shapes and once no graph + // holds them; the transforms of eager calls may still be in flight, so the + // device finishes them first. The destructor reports no error. ~MeshPlans() { int current = 0; cudaGetDevice(¤t); @@ -983,29 +1065,13 @@ struct MeshPlans { }; // The plans of the meshes of ``mesh`` with ``channels`` channels on the -// current device, kept across calls; beyond kPlanCacheSize shapes, the least -// recently used are released. The cache itself lives until the process ends -// and is never destroyed, since its plans would otherwise be released after -// the CUDA runtime shuts down. -constexpr size_t kPlanCacheSize = 16; - -std::shared_ptr mesh_plans(const FrameMesh& mesh, int channels) { - using Key = std::array; - using Cache = std::list>>; - static std::mutex cache_mutex; - static Cache& cache = *new Cache(); - const Key key = {c10::cuda::current_device(), mesh.n[0], mesh.n[1], mesh.n[2], - channels}; - // A plan that leaves the cache is released after the lock, since its - // release waits for the device. - std::shared_ptr evicted; - const std::lock_guard lock(cache_mutex); - for (auto entry = cache.begin(); entry != cache.end(); ++entry) { - if (entry->first == key) { - cache.splice(cache.begin(), cache, entry); - return entry->second; - } - } +// current device, with the assignment factors of B-splines of order +// ``order``, formed on ``stream``, which the creation waits for so that calls +// on every stream read them. +std::shared_ptr create_mesh_plans(const FrameMesh& mesh, + int channels, + int order, + cudaStream_t stream) { auto plans = std::make_shared(); plans->device = c10::cuda::current_device(); long long shape[3] = {mesh.n[0], mesh.n[1], mesh.n[2]}; @@ -1026,10 +1092,119 @@ std::shared_ptr mesh_plans(const FrameMesh& mesh, int channels) { 1, points, CUFFT_C2R, channels, &inverse_bytes), "plan creation"); plans->work_bytes = std::max(forward_bytes, inverse_bytes); - cache.emplace_front(key, plans); - if (cache.size() > kPlanCacheSize) { - evicted = std::move(cache.back().second); - cache.pop_back(); + const int64_t factor_count = + int64_t{mesh.n[0]} + mesh.n[1] + mesh.n[2] / 2 + 1; + plans->factor = + torch::empty({factor_count}, torch::TensorOptions() + .dtype(torch::kFloat32) + .device(torch::kCUDA, plans->device)); + const Float64Launch launch = float64_launch(factor_count); + assignment_factor<<>>( + mesh, order, integer_bspline_weights(order), + plans->factor.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + C10_CUDA_CHECK(cudaStreamSynchronize(stream)); + return plans; +} + +// Ends the hold of a graph on the plans it captured. CUDA runs it on a thread +// of its own, from which no CUDA call may be made, so the hold is a count +// that the cache reads rather than an owner that would release the plans. +void CUDART_CB release_graph_hold(void* plans) { + static_cast(plans)->graphs.fetch_sub(1); +} + +// Makes ``graph``, which a capture is recording, hold ``plans`` until it and +// every graph instantiated from it are destroyed. +void hold_in_graph(MeshPlans& plans, cudaGraph_t graph) { + cudaUserObject_t hold = nullptr; + C10_CUDA_CHECK(cudaUserObjectCreate(&hold, &plans, release_graph_hold, 1, + cudaUserObjectNoDestructorSync)); + plans.graphs.fetch_add(1); + const cudaError_t moved = + cudaGraphRetainUserObject(graph, hold, 1, cudaGraphUserObjectMove); + if (moved != cudaSuccess) { + // The reference stays with this call, whose release ends the hold. + cudaUserObjectRelease(hold, 1); + } + C10_CUDA_CHECK(moved); +} + +// Whether a CUDA graph captures the stream of a call, and that graph. +struct StreamCapture { + cudaStreamCaptureStatus status = cudaStreamCaptureStatusNone; + cudaGraph_t graph = nullptr; +}; + +StreamCapture stream_capture(cudaStream_t stream) { + StreamCapture capture; +#if CUDART_VERSION >= 12000 + C10_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &capture.status, nullptr, + &capture.graph)); +#else + C10_CUDA_CHECK(cudaStreamGetCaptureInfo_v2(stream, &capture.status, nullptr, + &capture.graph)); +#endif + return capture; +} + +// The plans of the meshes of ``mesh`` with ``channels`` channels and +// B-splines of order ``order`` on the current device, kept across calls. A +// call under ``capture`` makes the graph of the capture hold them and needs +// them created before, by an eager call, since their creation waits for the +// stream. The other calls release the least recently used plans that no +// graph holds past kPlanCacheSize shapes, which a call under capture may not +// do, since a release waits for the device. The cache itself lives until the +// process ends and is never destroyed, since its plans would otherwise be +// released after the CUDA runtime shuts down. +constexpr size_t kPlanCacheSize = 16; + +std::shared_ptr mesh_plans(const FrameMesh& mesh, + int channels, + int order, + const StreamCapture& capture, + cudaStream_t stream) { + using Key = std::array; + using Cache = std::list>>; + static std::mutex cache_mutex; + static Cache& cache = *new Cache(); + const Key key = {c10::cuda::current_device(), + mesh.n[0], + mesh.n[1], + mesh.n[2], + channels, + order}; + // Plans that leave the cache are released after the lock, since their + // release waits for the device. + std::vector> evicted; + const std::lock_guard lock(cache_mutex); + const auto found = + std::find_if(cache.begin(), cache.end(), + [&](const auto& entry) { return entry.first == key; }); + if (found != cache.end()) { + cache.splice(cache.begin(), cache, found); + } else { + TORCH_CHECK(capture.status == cudaStreamCaptureStatusNone, + "latent_charge_sum: a CUDA graph captures only meshes that an " + "eager call has evaluated, and none has evaluated the mesh of ", + mesh.n[0], " x ", mesh.n[1], " x ", mesh.n[2], + " points, channel count ", channels, " and B-spline order ", + order); + cache.emplace_front(key, create_mesh_plans(mesh, channels, order, stream)); + } + std::shared_ptr plans = cache.front().second; + if (capture.status == cudaStreamCaptureStatusActive) { + hold_in_graph(*plans, capture.graph); + } else if (capture.status == cudaStreamCaptureStatusNone) { + // From the least recently used back to, but short of, the front entry, + // which holds the plans of this call. + auto entry = cache.end(); + while (cache.size() > kPlanCacheSize && --entry != cache.begin()) { + if (entry->second->graphs.load() == 0) { + evicted.push_back(std::move(entry->second)); + entry = cache.erase(entry); + } + } } return plans; } @@ -1135,6 +1310,14 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, "latent_charge_sum: owners other than every node lie in one " "frame"); const auto stream = at::cuda::getCurrentCUDAStream(); + // The device reads pinned cells in place, and a copy of pageable ones; a + // CUDA graph that captures the call reads the cells again at every replay, + // which only pinned memory allows. + const StreamCapture capture = stream_capture(stream); + const bool pinned = box.is_pinned(); + TORCH_CHECK(capture.status == cudaStreamCaptureStatusNone || pinned, + "latent_charge_sum: a CUDA graph that captures periodic frames " + "reads their cells at every replay, from box in pinned memory"); // === Step 1. The meshes of the frames, sized on the host === std::vector frames(frame_count); @@ -1142,33 +1325,20 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, const double* cells = box.data_ptr(); int64_t mesh_total = 0; int64_t spectrum_size = 0; - int64_t factor_size = 0; + int mode_blocks = 1; size_t work_bytes = 0; for (int f = 0; f < frame_count; ++f) { - frames[f] = frame_mesh(cells + 9 * f, spacing); + frames[f] = frame_mesh_size(cells + 9 * f, spacing); frames[f].offset = mesh_total; mesh_total += channels * mesh_points(frames[f]); const int half = frames[f].n[2] / 2 + 1; spectrum_size = std::max(spectrum_size, int64_t{channels} * frames[f].n[0] * frames[f].n[1] * half); - factor_size = - std::max(factor_size, int64_t{frames[f].n[0]} + frames[f].n[1] + half); - plans[f] = mesh_plans(frames[f], channels); + mode_blocks = std::max(mode_blocks, influence_blocks(frames[f])); + plans[f] = mesh_plans(frames[f], channels, static_cast(order), capture, + stream); work_bytes = std::max(work_bytes, plans[f]->work_bytes); } - const auto byte_options = torch::TensorOptions().dtype(torch::kUInt8); - auto host_frames = - torch::empty({static_cast(sizeof(FrameMesh)) * frame_count}, - byte_options.pinned_memory(true)); - std::memcpy(host_frames.data_ptr(), frames.data(), - sizeof(FrameMesh) * frame_count); - const IntegerWeights integer_weight = - integer_bspline_weights(static_cast(order)); - // A few blocks per multiprocessor fill the device; every thread of the - // influence function then walks several modes, which keeps the block sums - // few. - const int mode_blocks = - 4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount; // === Step 2. The buffers of the call === const int64_t average_nodes = (node_count + frame_count - 1) / frame_count; @@ -1186,6 +1356,8 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, Workspace workspace; const int64_t frames_at = workspace.reserve(sizeof(FrameMesh) * frame_count); + const int64_t cells_at = + workspace.reserve(pinned ? 0 : int64_t{frame_count} * 9); const int64_t frame_ptr_at = workspace.reserve(frame_count + 1); const int64_t slice_charge_at = workspace.reserve(int64_t{frame_count} * slice_count * 2); @@ -1201,7 +1373,6 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, const int64_t mesh_at = workspace.reserve(mesh_total); const int64_t spectrum_at = workspace.reserve(spectrum_size); const int64_t work_at = workspace.reserve(work_bytes); - const int64_t factor_at = workspace.reserve(factor_size); const int64_t block_sum_at = workspace.reserve(int64_t{mode_blocks} * kModeSums); const int64_t frame_virial_at = @@ -1223,17 +1394,35 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, float* potential_mesh = charge_mesh; auto* spectrum = workspace.at(spectrum_at); void* work = workspace.at(work_at); - auto* factor = workspace.at(factor_at); auto* block_sum = workspace.at(block_sum_at); auto* frame_virial = workspace.at(frame_virial_at); auto* potential = workspace.at(potential_at); auto* owner_slice = workspace.at(owner_slice_at); auto* owner_local = workspace.at(owner_local_at); - // The copy from the pinned staging buffer records its completion, before - // which the allocator does not hand the buffer out again. - torch::from_blob(frames_ptr, {host_frames.numel()}, - byte_options.device(coord.device())) - .copy_(host_frames, /*non_blocking=*/true); + const double* device_cells = nullptr; + if (pinned) { + void* mapped = nullptr; + C10_CUDA_CHECK( + cudaHostGetDevicePointer(&mapped, const_cast(cells), 0)); + device_cells = static_cast(mapped); + } else { + // Pageable cells are staged in pinned memory, from which the copy runs + // without waiting for the stream; the host allocator keeps the staging + // memory until the copy has completed. + auto staged = + torch::empty({frame_count, 9}, box.options().pinned_memory(true)); + std::memcpy(staged.data_ptr(), cells, + sizeof(double) * 9 * frame_count); + auto* copy = workspace.at(cells_at); + torch::from_blob(copy, {frame_count, 9}, coord.options()) + .copy_(staged, /*non_blocking=*/true); + device_cells = copy; + } + for (int f = 0; f < frame_count; ++f) { + store_frame_mesh<<<1, kGeometryThreads, 0, stream>>>( + frames[f], device_cells + 9 * f, prefactor, frames_ptr + f); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } C10_CUDA_CHECK( cudaMemsetAsync(fixed_mesh, 0, sizeof(int64_t) * mesh_total, stream)); @@ -1254,11 +1443,13 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, dispatch_order(order, [&](auto order_constant) { constexpr int kOrder = decltype(order_constant)::value; auto* node_stencil = workspace.at>(stencils_at); - node_stencils<<>>( - coord.data_ptr(), charge.data_ptr(), - node_mask.data_ptr(), frame_ptr, frames_ptr, frame_stat, - frame_count, node_count, channels, static_cast(channel), - node_stencil, node_charge); + const Float64Launch stencil_launch = float64_launch(3 * node_count); + node_stencils + <<>>( + coord.data_ptr(), charge.data_ptr(), + node_mask.data_ptr(), frame_ptr, frames_ptr, frame_stat, + frame_count, node_count, channels, static_cast(channel), + node_stencil, node_charge); C10_CUDA_KERNEL_LAUNCH_CHECK(); // === Step 4. The charge meshes === spread_charges @@ -1283,17 +1474,13 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, reinterpret_cast(spectrum)), "transform"); } - const int half = grid.n[2] / 2 + 1; - const int factor_count = grid.n[0] + grid.n[1] + half; - assignment_factor<<>>( - grid, static_cast(order), integer_weight, factor); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - apply_influence<<>>( - grid, channels, static_cast(kernel), - coefficient.data_ptr(), gaussians, prefactor, factor, spectrum, - block_sum); + const int blocks = influence_blocks(grid); + apply_influence<<>>( + frames_ptr + f, channels, static_cast(kernel), + coefficient.data_ptr(), gaussians, + plan.factor.data_ptr(), spectrum, block_sum); C10_CUDA_KERNEL_LAUNCH_CHECK(); - mode_virial<<<1, kThreads, 0, stream>>>(block_sum, mode_blocks, + mode_virial<<<1, kThreads, 0, stream>>>(block_sum, blocks, frame_virial + 9 * f); C10_CUDA_KERNEL_LAUNCH_CHECK(); { @@ -1358,8 +1545,9 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, auto charge_gradient = torch::empty({owner_count, channels}, float_options); auto atom_virial = torch::empty({owner_count, 9}, float_options); auto virial = torch::empty({frame_count, 9}, float_options); - finish_owners<<(owner_count, 9 * frame_count)), - kThreads, 0, stream>>>( + const Float64Launch owner_launch = + float64_launch(std::max(owner_count, 9 * frame_count)); + finish_owners<<>>( potential, force.data_ptr(), owner_coord, owner_charge, owner_mask, owner_frame_ptr, frame_virial, owner_local, owner_total, frame_stat, frame_count, owner_count, channels, static_cast(channel), diff --git a/source/op/pt/long_range/ops.cc b/source/op/pt/long_range/ops.cc index 364fa6c2c2..236e61c147 100644 --- a/source/op/pt/long_range/ops.cc +++ b/source/op/pt/long_range/ops.cc @@ -60,8 +60,18 @@ // (F, 9). The mesh of a periodic frame holds along every lattice vector a_d // the smallest even number of points without a prime factor above 7 of at // least ceil(|a_d| / spacing), and spreads the charges with B-splines of -// order `order`. The force and the virial hold the charges fixed; the response -// of the charges is the caller's, through the charge gradient. +// order `order`. The operator sizes the meshes from `box` on the host; the +// device reads pinned cells in place, while the stream runs the call, and a +// copy of pageable ones. A CUDA graph that captures a call reads the cells +// again at every replay, so it needs `box` in pinned memory and evaluates the +// cells there at replay, provided their meshes keep the sizes of the capture. +// It captures only meshes whose size, with the channel count and the order, +// an eager call has evaluated: an eager call forms the cuFFT plans and the +// assignment factors of a mesh size, which the operator keeps for the sizes +// used last and for as long as a graph holds them. The caller keeps pinned +// cells alive and unchanged until the stream has passed a call that reads +// them. The force and the virial hold the charges fixed; the response of the +// charges is the caller's, through the charge gradient. // // `communicator` is an int64 CPU tensor of one element holding the address of // the caller's MPI_Comm handle, or zero. A handle of several ranks decomposes diff --git a/source/tests/infer/gen_dpa4c_long_range.py b/source/tests/infer/gen_dpa4c_long_range.py index 8fa069eb05..bb33fb39c8 100644 --- a/source/tests/infer/gen_dpa4c_long_range.py +++ b/source/tests/infer/gen_dpa4c_long_range.py @@ -160,6 +160,9 @@ def _check_metadata(pt2_path: str, lower: str, periodic: bool = False) -> None: # carries the rebuild of every constant that holds it. assert md["dim_chg_spin"] == 0 assert len(md["charge_state_constants"]) == 5 + # The fused long-range operator serves the charge network of one + # channel and reads the source CSR for the per-atom virial alone. + assert md["canonical_requires_source_csr"] is False def _write_ref(pt2_path: str, ref_path: str, periodic: bool = False) -> None: diff --git a/source/tests/pt_expt/kernels/test_long_range_sum.py b/source/tests/pt_expt/kernels/test_long_range_sum.py index b61c7390c9..165292f733 100644 --- a/source/tests/pt_expt/kernels/test_long_range_sum.py +++ b/source/tests/pt_expt/kernels/test_long_range_sum.py @@ -144,6 +144,9 @@ def _evaluate( ([37, 1, 0, 300], 1, 5, -1), ([2048], 3, 0, -1), ([5, 129], 8, 3, 2), + # A frame of 47 owner blocks of 24 chunks, which takes the wide tiles + # of the pair kernel on devices of up to 3384 resident block slots. + ([6000], 2, 4, 1), ], ) def test_matches_float64_pair_sum( @@ -189,10 +192,21 @@ def _assert_close( assert error <= tolerance * scale, (name, error, scale) -@pytest.mark.parametrize("kernel_name", ["les", "sog"]) -def test_invariants(kernel_name: str) -> None: +@pytest.mark.parametrize( + ("kernel_name", "n_node"), + [ + ("les", [64, 200]), + ("sog", [64, 200]), + # A first frame that ends inside a block of the pair kernel, whose + # chunks then reach past the partners of that frame. The erf kernel + # decides its far field over the lanes of a warp, which such a block + # shares between the frames, so that it is checked on frames of whole + # warps. + ("sog", [150, 200]), + ], +) +def test_invariants(kernel_name: str, n_node: list[int]) -> None: """Frames are independent, forces cancel, and the virial is the pair virial.""" - n_node = [64, 200] system = _system(n_node, 2, 0, 11, 0.0) coord, charge, node_mask, n_node_t, total_charge = system together = latent_charge_sum( @@ -206,7 +220,7 @@ def test_invariants(kernel_name: str) -> None: COULOMB_CONSTANT, ) _, _, force, atom_virial, virial = together - # Each frame alone gives the same result. + # The first frame alone gives the same result. first = n_node[0] alone = latent_charge_sum( coord[:first], @@ -438,6 +452,96 @@ def test_mesh_frames_are_independent() -> None: torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0) +def _mesh_sizes(box: torch.Tensor, spacing: float) -> list[tuple[int, int, int]]: + """The mesh size of every frame of ``box``.""" + return [mesh_dims_of_cells(cell.reshape(1, 9).numpy(), spacing) for cell in box] + + +def test_captured_meshes_read_the_cells_at_replay() -> None: + """A CUDA graph of periodic frames evaluates the cells of its box at every replay. + + The graph reads the cells where they lie in pinned memory, so that a + replay after a strain that keeps the sizes of the meshes agrees bit for + bit with an eager evaluation of the strained cells; cells in pageable + memory cannot be captured. + """ + system = _system([64, 200], 2, 0, 11, 0.0) + kernel = KERNELS["sog"]().to("cuda") + + def evaluate(box: torch.Tensor) -> tuple[torch.Tensor, ...]: + return latent_charge_sum( + *system, kernel, 1, COULOMB_CONSTANT, box=box, spacing=0.5 + ) + + cells = _cells(2, 9.0) + strained = cells * 1.005 + assert _mesh_sizes(strained, 0.5) == _mesh_sizes(cells, 0.5) + expected = {"first": evaluate(cells), "strained": evaluate(strained)} + with pytest.raises(RuntimeError, match="pinned"): + with torch.cuda.graph(torch.cuda.CUDAGraph()): + evaluate(cells) + box = cells.pin_memory() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = evaluate(box) + for name, cell in (("strained", strained), ("first", cells)): + box.copy_(cell) + graph.replay() + torch.cuda.synchronize() + for value, reference in zip(captured, expected[name], strict=True): + torch.testing.assert_close( + value, reference, rtol=0, atol=0, msg=lambda m, n=name: f"{n}: {m}" + ) + + +def test_captured_meshes_outlive_other_shapes() -> None: + """A CUDA graph of periodic frames replays after 17 other mesh shapes, bit for bit. + + The operator keeps the transforms of 16 mesh shapes and releases the least + recently used past them, except those of a live graph, which runs them at + every replay. + """ + system = _system([64, 200], 2, 0, 11, 0.0) + kernel = KERNELS["sog"]().to("cuda") + + def evaluate(box: torch.Tensor) -> tuple[torch.Tensor, ...]: + return latent_charge_sum( + *system, kernel, 1, COULOMB_CONSTANT, box=box, spacing=0.5 + ) + + box = _cells(2, 9.0).pin_memory() + expected = evaluate(box) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = evaluate(box) + others: set[tuple[int, int, int]] = set() + length = 9.0 + while len(others) < 17: + length += 0.5 + cells = _cells(2, length) + evaluate(cells) + others |= set(_mesh_sizes(cells, 0.5)) - set(_mesh_sizes(box, 0.5)) + graph.replay() + for value, reference in zip(captured, expected, strict=True): + torch.testing.assert_close(value, reference, rtol=0, atol=0) + + +def test_a_capture_needs_an_eager_call_for_its_meshes() -> None: + """A CUDA graph captures only meshes that an eager call has evaluated.""" + system = _system([12], 1, 0, 9, 0.0) + kernel = KERNELS["les"]().to("cuda") + # A cell whose mesh of 16 x 16 x 240 points no other test evaluates. + box = torch.tensor( + [[4.0, 0.0, 0.0, 0.0, 4.0, 0.0, 0.0, 0.0, 60.0]], dtype=torch.float64 + ).pin_memory() + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="none has evaluated the mesh of 16 x 16"): + with torch.cuda.graph(graph): + latent_charge_sum( + *system, kernel, 0, COULOMB_CONSTANT, box=box, spacing=0.25 + ) + + def test_mixed_frames_are_rejected() -> None: """A call holds periodic frames or non-periodic ones, not both.""" system = _system([10, 10], 1, 0, 1, 0.0) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index c48dc0ae4c..0c69d53883 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -961,6 +961,13 @@ def test_compressed_frozen_model_matches_uncompressed( dp = DeepPot(str(frozen), auto_batch_size=False) expected_lower = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph" assert dp.deep_eval.metadata["lower_input_kind"] == expected_lower + if expected_lower == "dpa4c_canonical": + # The fused long-range operator serves the charge network of one + # channel and reads the source CSR for the per-atom virial alone; the + # composition of operators that serves three channels reads it always. + assert dp.deep_eval.metadata["canonical_requires_source_csr"] == ( + CHANNELS[kernel] > 1 + ) rng = np.random.default_rng(3) coord = rng.uniform(0.0, 7.0, (1, 40, 3)) + 100.0 atype = rng.integers(0, 2, (1, 40)) @@ -1515,6 +1522,10 @@ def test_periodic_compressed_frozen_model_matches_uncompressed( dp = DeepPot(str(frozen), auto_batch_size=False) expected_lower = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph" assert dp.deep_eval.metadata["lower_input_kind"] == expected_lower + if expected_lower == "dpa4c_canonical": + assert dp.deep_eval.metadata["canonical_requires_source_csr"] == ( + CHANNELS[kernel] > 1 + ) names = ("energy", "force", "virial", "atom_energy", "atom_virial") rng = np.random.default_rng(3) box = np.array([[[12.0, 0.0, 0.0], [0.8, 11.0, 0.0], [-0.5, 0.6, 13.0]]]) @@ -1546,6 +1557,212 @@ def test_periodic_compressed_frozen_model_matches_uncompressed( dp.eval(periodic, np.zeros((1, 9)), atype) +def _canonical_inputs( + model: torch.nn.Module, + coord: np.ndarray, + atype: np.ndarray, + box: np.ndarray | None, +) -> tuple: + """The compact canonical graph of the frames, as one rank builds it, and the node inputs of the module.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + build_neighbor_graph, + ) + from deepmd.pt_expt.utils.canonical_graph import ( + canonical_graph_from_neighbor_graph, + ) + + nf = atype.shape[0] + cells = None if box is None else box.reshape(nf, 9) + graph = build_neighbor_graph( + coord, atype, cells, model.get_rcut(), canonicalize=True + ) + n_node = torch.as_tensor(graph.n_node, dtype=torch.int64, device=env.DEVICE) + graph = NeighborGraph( + n_node=n_node, + n_local=n_node, + edge_index=torch.as_tensor(graph.edge_index, device=env.DEVICE), + edge_vec=torch.as_tensor( + graph.edge_vec, dtype=torch.float32, device=env.DEVICE + ), + edge_mask=torch.as_tensor(graph.edge_mask, device=env.DEVICE), + destination_order=torch.as_tensor(graph.destination_order, device=env.DEVICE), + destination_row_ptr=torch.as_tensor( + graph.destination_row_ptr, dtype=torch.int64, device=env.DEVICE + ), + source_order=torch.as_tensor(graph.source_order, device=env.DEVICE), + source_row_ptr=torch.as_tensor( + graph.source_row_ptr, dtype=torch.int64, device=env.DEVICE + ), + destination_sorted=bool(graph.destination_sorted), + ) + compact = canonical_graph_from_neighbor_graph(graph) + return ( + _tensor(atype.reshape(-1), torch.int64), + compact, + _tensor(coord.reshape(-1, 3)), + None if cells is None else torch.as_tensor(cells, dtype=torch.float64), + ) + + +# The layouts of the fused parity test: the width of the energy fitting, which +# selects the per-layer kernels at 32 and the fused network at 96, the hidden +# layers of the charge network, and the frames of the batch. +FUSED_LAYOUTS = { + "per_layer": {"width": 32, "charge": [16, 16], "frames": 2}, + "fused": {"width": 96, "charge": [16, 16, 16], "frames": 1}, +} + + +@pytest.mark.skipif( + env.DEVICE.type != "cuda", reason="the compact canonical lower is CUDA only" +) +@pytest.mark.parametrize("tile", [0, 16]) +@pytest.mark.parametrize("layout", ["per_layer", "fused"]) +@pytest.mark.parametrize("periodic", [False, True]) +@pytest.mark.parametrize(("kernel", "source"), [("les", "neutral"), ("sog", None)]) +def test_fused_operator_matches_composition( + monkeypatch, + kernel: str, + source: str | None, + periodic: bool, + layout: str, + tile: int, +) -> None: + """The fused long-range operator evaluates a compressed model like the composition. + + With a charge network of one channel, the compact canonical lower + evaluates the descriptor, both networks, the long range and the force + assembly in one operator; the composition of operators, which serves the + other charge networks, is the reference. The energy fitting takes the + per-layer kernels or the fused network, over the whole node axis or in + runs of 16 nodes with a shorter last one, with and without the per-atom + virial, on two frames, the second with fewer real atoms, or on one. + """ + from deepmd.pt_expt.kernels.dpa4c import ( + canonical_long_range, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_operator_arguments, + ) + + shape = FUSED_LAYOUTS[layout] + monkeypatch.setenv("DP_NODE_TILE", str(tile)) + config = _config(kernel, constraint=None if source is None else {"source": source}) + config["descriptor"].update(n_radial=8, precision="float32") + config["fitting_net"].update( + neuron=[shape["width"]] * 2, activation_function="silu", resnet_dt=False + ) + config["fitting_net"]["precision"] = "float32" + config["long_range"].update( + channels=1, + charge_net={"neuron": shape["charge"], "seed": 23}, + periodic=periodic, + ) + if kernel == "sog": + config["long_range"]["kernel"] = {"type": "sog", "n_gaussian": 4} + model = get_model(config).to(env.DEVICE).eval() + generator = torch.Generator().manual_seed(5) + with torch.no_grad(): + for parameter in model.parameters(): + noise = torch.randn(parameter.shape, generator=generator) + parameter.add_(0.05 * noise.to(parameter)) + if periodic: + model.atomic_model.long_range.use_mesh_evaluation() + model.atomic_model.enable_compression(0.5) + assert canonical_long_range.fused_long_range_eligible(model.atomic_model.long_range) + frames = shape["frames"] + rng = np.random.default_rng(3) + box = np.array([[[12.0, 0.0, 0.0], [0.8, 11.0, 0.0], [-0.5, 0.6, 13.0]]] * frames) + coord = ( + rng.uniform(0.0, 1.0, (frames, 30, 3)) @ box[0] + if periodic + else (rng.uniform(0.0, 7.0, (frames, 30, 3)) + 100.0) + ) + atype = rng.integers(0, 2, (frames, 30)) + atype[-1, 24:] = -1 + atype_t, graph, coord_t, box_t = _canonical_inputs( + model, coord, atype, box if periodic else None + ) + # The rules of the fitting routes take the network of the layout on every + # run of the energy fitting; a device on which they send the runs of the + # fused layout to the per-layer kernels has no fused network to compare. + weights = fitting_operator_arguments(model.atomic_model.fitting_net).weights + nodes = atype_t.shape[0] + runs = {nodes} if tile == 0 else {tile, nodes % tile or tile} + routes = [int(torch.ops.deepmd.graph_fitting_route_rule(weights, n)) for n in runs] + if layout == "fused" and -1 in routes: + pytest.skip("the fitting rules take the per-layer kernels on this device") + assert all((route == -1) == (layout == "per_layer") for route in routes), routes + + def evaluate(do_atomic_virial: bool) -> dict[str, torch.Tensor]: + return model.forward_lower_canonical_graph( + atype_t, + graph.n_node, + graph.n_local, + graph.source, + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + do_atomic_virial=do_atomic_virial, + coord=coord_t, + box=box_t, + ) + + for do_atomic_virial in (True, False): + fused = evaluate(do_atomic_virial) + with monkeypatch.context() as patch: + patch.setattr( + canonical_long_range, "fused_long_range_eligible", lambda _: False + ) + composed = evaluate(do_atomic_virial) + assert set(fused) == set(composed) + for key, value in composed.items(): + # Both evaluate in float32 and sum in other orders, so they agree + # to a few roundings of the largest entry of each output; the + # long range dominates these outputs. + expected = value.detach().to(torch.float64) + error = (fused[key].detach().to(torch.float64) - expected).abs().max() + scale = expected.abs().max() + assert error <= 1e-5 * scale + 1e-6, (key, error.item(), scale.item()) + + +def test_library_without_the_fused_operator_keeps_the_composition(monkeypatch) -> None: + """A library without the fused long-range operator evaluates the composition. + + The fused operator is optional beside the short-range suite: without it, + a module that it would serve takes the composition of operators, and the + archive records that its compact lower reads the source CSR. + """ + from deepmd.pt_expt.kernels.dpa4c import ( + canonical, + canonical_long_range, + ) + from deepmd.pt_expt.utils.serialization import ( + _canonical_requires_source_csr, + ) + + config = _config("les") + config["descriptor"]["precision"] = "float32" + config["fitting_net"].update( + neuron=[32, 32], activation_function="silu", resnet_dt=False + ) + config["fitting_net"]["precision"] = "float32" + model = get_model(config) + monkeypatch.setattr( + canonical, + "operator_available", + lambda name: name != "dpa4c_canonical_compress_long_range_energy_force", + ) + assert canonical.op_available() + assert not canonical.long_range_op_available() + assert not canonical_long_range.fused_long_range_eligible( + model.atomic_model.long_range + ) + assert _canonical_requires_source_csr(model, "dpa4c_canonical") + + # --------------------------------------------------------------------------- # Periodic fixed k-grid mode # --------------------------------------------------------------------------- From 0289a0d33dede339463aee37fa4de953aa75ee43 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 25 Oct 2026 21:04:48 +0800 Subject: [PATCH 03/15] perf(long-range): fuse the charge network and harden the LAMMPS path The charge network of a module that the fused long-range operator evaluates runs through the fused fitting network or the per-layer kernels by a fitting-route table of its own, which the freeze measures like that of the energy fitting and passes as the new optional argument charge_fitting_routes. A charge depends on the descriptor of its own node alone, so its cotangent is the derivative of the energy with respect to the charge times its Jacobian row, the input cotangent of a unit seed. A run of the fused network forms the charges and those rows in one pass before the long range; a run of the per-layer kernels forms its charges there and its rows after the energy fitting, from pre-activations kept for every node, which on a V100 at 512 atoms shortens the LAMMPS step by 3.3% against forming them before the long range. One elementwise pass adds the scaled rows to the cotangent of the energy fitting. The fused network is compiled for the hidden width of 64, that of the default charge network: 32-node panels on clusters of two CTAs or on one CTA for small runs, and 32- or 64-node panels for large ones. In LAMMPS with Kokkos on compute capability 12.0, the kernels of the charge network of the trained Air model take 39 us per step at 10,296 atoms against 47 us, and 429 against 477 to 493 us at 100,800 atoms (Nsight Systems, periodic systems and clusters with LES and SOG modules); the forces move by up to 3.3e-6 of the largest, the rounding of the Jacobian product, and models without a module keep their bits. The pair kernel of the latent-charge direct sum takes its tile and its partner chunks from the resident blocks of a multiprocessor, which it now queries for each device of a process rather than once per process, so that a process that drives devices of different kinds launches each with its own occupancy. The limits and kernel codes of the operator, its handling of a call without nodes or frames and the slices of its per-frame sums are defined once for the direct and the mesh sum, and the domain query returns the rank count it is used for. The two fused DPA4C energy-force operators share the plan of the fitting runs (the routes of the full and the tail runs and their buffers) and the epilogue that forms the node results, the frame totals and the per-atom virial, whose optional magnetic, source-CSR and long-range terms are passed as groups; the outputs of a model without a long-range module are unchanged. pair_style deepmd rejects at init an ensemble on a non-periodic system whose first model carries no long-range module, which the steps without a model deviation evaluate alone without a cell. deepmd/kk passes the positions relative to the origin of the box only to a model with a long-range module, and keeps the positions of LAMMPS in place for the others. The pt backend rejects a long-range module of a model with a tabulated pair potential, as it does for the other model types. The default of the source-CSR record of a compact lower without one is read in one expression. The C++ test runner waits for the canonical and the long-range fixture generators separately. test_lammps_dpa4c_long_range_kokkos.py, run by test_lammps_kokkos.sh, evaluates the compressed long-range fixtures under deepmd/kk in a box away from the origin, with and without a cell and a charge state, on one and two ranks, against their Python references; without the shift of the origin its periodic cases fail. The MPI runner of the pair tests runs deepmd/kk on request, and a bounded run ends its ranks with the launcher. New tests check that every route of either network of the fused operator returns the bits of the rules, launches the kernels it names and leaves the other network at its rules, that the table the route tuner holds for a charge network, which the energy model passes, routes it, and that the tuner measures the charge network apart. Comments and the long-range documentation are corrected where they had drifted from the code. --- .../utils/neighbor_graph/derivatives.py | 4 +- .../utils/neighbor_graph/pair_graph.py | 8 +- .../dpmodel/utils/neighbor_graph/topology.py | 3 +- deepmd/pt/model/model/__init__.py | 22 +- .../kernels/dpa4c/canonical_long_range.py | 18 +- deepmd/pt_expt/kernels/fitting_routes.py | 121 ++-- deepmd/pt_expt/model/ener_model.py | 7 +- doc/model/long-range.md | 8 +- source/api_cc/include/DeepPot.h | 6 +- source/api_cc/src/DeepPotPTExpt.cc | 10 +- source/api_cc/src/longRangePTExpt.h | 4 +- source/install/test_cc_local.sh | 4 +- source/install/test_lammps_kokkos.sh | 7 + source/lmp/pair_deepmd.cpp | 12 + source/lmp/pair_deepmd_kokkos.cpp | 16 +- source/lmp/pair_deepmd_kokkos.h | 6 +- source/lmp/tests/lammps_test_utils.py | 33 +- .../lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py | 23 +- .../test_lammps_dpa4c_long_range_kokkos.py | 144 +++++ .../tests/test_lammps_dpa4c_long_range_pt2.py | 8 +- source/op/pt/CMakeLists.txt | 6 +- source/op/pt/dpa4c/graph_compress.cu | 580 +++++++++++------- source/op/pt/dpa4c/ops.cc | 2 +- source/op/pt/graph_fitting.cu | 39 +- source/op/pt/long_range/charge_sum.h | 15 +- source/op/pt/long_range/direct_sum.cu | 91 ++- source/op/pt/long_range/domain_exchange.cc | 12 +- source/op/pt/long_range/domain_exchange.h | 10 +- source/op/pt/long_range/frame_sums.cuh | 26 + source/op/pt/long_range/latent_charge_sum.cu | 28 +- source/op/pt/long_range/mesh_sum.cu | 42 +- source/tests/infer/gen_dpa4c_long_range.py | 4 +- .../pt_expt/descriptor/test_dpa1_cuda.py | 2 +- .../pt_expt/descriptor/test_dpa4c_cuda.py | 2 +- source/tests/pt_expt/model/test_long_range.py | 393 +++++++++++- 35 files changed, 1252 insertions(+), 464 deletions(-) create mode 100644 source/lmp/tests/test_lammps_dpa4c_long_range_kokkos.py diff --git a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py index 5e3f34ec04..7f643bd418 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -7,7 +7,9 @@ Conventions (see the unified edge-nlist design discussion, wanghan-iapcm/deepmd-kit#4): edge_vec_e = r_src - r_dst ; F_k = sum_{dst=k} g - sum_{src=k} g per-edge virial w_e = -g_e (x) edge_vec_e - atom virial attributed FULL-TO-src (canonical TF==pt-legacy convention) + atom virial attributed FULL-TO-src (canonical TF==pt-legacy convention); + a pair set (pair_force_virial) splits its pair virial evenly between the + two endpoints instead per-frame virial = sum over the edges of that frame of w_e (DeePMD virials are per frame; a multi-frame NeighborGraph must NOT collapse frames) Padding/guard edges (edge_mask == 0) are zeroed before any scatter. diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py index e474abd4ed..aa0bd0ef39 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -189,10 +189,10 @@ def build_pair_graph( def empty_pair_graph(like: Array) -> PairGraph: """The pair topology of a batch without non-periodic frames: guard pairs only. - The guard pairs number :data:`_GUARD_PAIRS`, unlike the two rows of the - index and the three components of the vectors: a compiled graph that met - a pair axis of the length of one of those would identify the two, and - split at the reductions over the pairs. + The guard pairs number :data:`_GUARD_PAIRS`, which differs from the two + rows of the index and the three components of the vectors: a compiled + graph would identify a pair axis of the length of one of those axes with + it, and split at the reductions over the pairs. """ xp = array_api_compat.array_namespace(like) device = array_api_compat.device(like) diff --git a/deepmd/dpmodel/utils/neighbor_graph/topology.py b/deepmd/dpmodel/utils/neighbor_graph/topology.py index 9a504fc7f2..e03e9aa2cf 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/topology.py +++ b/deepmd/dpmodel/utils/neighbor_graph/topology.py @@ -146,7 +146,8 @@ def build_topologies( If cells are given for a declaration of the pairs alone, or a pair exclusion for a periodic batch, whose reciprocal-space sum is global. ValueError - If a name is unknown, or a build parameter is missing. + If a name is unknown, a build parameter is missing, or a cell spans + no lattice. """ names = tuple(names) unknown = [name for name in names if name not in TOPOLOGY_TYPES] diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index b65604c40c..42894de746 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -108,6 +108,15 @@ ) +def _reject_long_range(model_params: dict) -> None: + """Reject a long-range module, which only the pt_expt backend serves.""" + if model_params.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) + + def _get_standard_model_components(model_params: dict, ntypes: int) -> tuple: if "type_embedding" in model_params: raise ValueError( @@ -208,11 +217,7 @@ def get_linear_model(model_params: dict) -> BaseModel: "sub-model: add an `inner_potential` sub-model to the " "composition instead." ) - if sub.get("long_range") is not None: - raise NotImplementedError( - "long-range modules (`long_range`) are supported by the pt_expt " - "backend only" - ) + _reject_long_range(sub) if any( sub.get("type") == "inner_potential" for sub in model_params.get("models", []) ): @@ -378,6 +383,7 @@ def _get_bridged_linear_model(model_params: dict) -> BaseModel: def get_zbl_model(model_params: dict) -> DPZBLModel: + _reject_long_range(model_params) model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) descriptor, fitting, _ = _get_standard_model_components(model_params, ntypes) @@ -419,11 +425,7 @@ def get_standard_model(model_params: dict) -> BaseModel: "expands the flag), or spell the composition explicitly with " '`type: "linear_ener"` and an `inner_potential` sub-model.' ) - if model_params.get("long_range") is not None: - raise NotImplementedError( - "long-range modules (`long_range`) are supported by the pt_expt " - "backend only" - ) + _reject_long_range(model_params) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py index e7bd81b80f..b588eb04ef 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -48,6 +48,7 @@ def dpa4c_canonical_compress_long_range_energy_force( box: torch.Tensor | None, atom_bias: torch.Tensor, fitting_routes: list[int] | None = None, + charge_fitting_routes: list[int] | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -97,7 +98,11 @@ def dpa4c_canonical_compress_long_range_energy_force( Combined atomic energy bias with shape ``(ntypes,)`` in eV. fitting_routes Fitting-route table of the energy fitting, as for the fused operator - without a long-range module, or ``None`` for the measured table. + without a long-range module. + charge_fitting_routes + Fitting-route table of the charge network. ``None`` or an empty table + leaves the routes of its network to the rules of the operator, for + either table. Returns ------- @@ -124,7 +129,6 @@ def dpa4c_canonical_compress_long_range_energy_force( from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( compressed_operator_arguments, ) - from deepmd.pt_expt.kernels.fitting_routes import fitting_routes as measured_routes from deepmd.pt_expt.kernels.graph_fitting import ( ensure_registered as ensure_fitting_registered, ) @@ -145,7 +149,6 @@ def dpa4c_canonical_compress_long_range_energy_force( kernel, coefficient = kernel_operator_arguments(long_range.kernel) if communicator is None: communicator = torch.zeros(1, dtype=torch.int64, device="cpu") - routes = measured_routes(fitting) if fitting_routes is None else fitting_routes return torch.ops.deepmd.dpa4c_canonical_compress_long_range_energy_force( graph.edge_vec, graph.source, @@ -188,7 +191,8 @@ def dpa4c_canonical_compress_long_range_energy_force( # No table is passed as None rather than as an empty list: the # AOTInductor proxy executor of torch 2.13 and earlier hands a # serialized empty int list to the operator as None. - list(routes) if routes else None, + list(fitting_routes) if fitting_routes else None, + list(charge_fitting_routes) if charge_fitting_routes else None, ) @@ -223,8 +227,10 @@ def dpa4c_canonical_long_range_energy_force( long_range Long-range module whose ``operator_eligible()`` holds. graph - Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` and - whose ``destination_row_ptr`` has shape ``(N + 1,)`` in ``int64``. + Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` or + ``int64``, or shape ``(E, W)`` in ``uint8`` for the exported bytes of + those words, and whose ``destination_row_ptr`` has shape ``(N + 1,)`` + in ``int64``. Its source CSR is complete. atype Flat node atom types with shape ``(N,)``. coord diff --git a/deepmd/pt_expt/kernels/fitting_routes.py b/deepmd/pt_expt/kernels/fitting_routes.py index d1dab441c5..354668107b 100644 --- a/deepmd/pt_expt/kernels/fitting_routes.py +++ b/deepmd/pt_expt/kernels/fitting_routes.py @@ -53,6 +53,12 @@ that of the alternating replays as the device idles more. On these compositions a route that the measurement selects is therefore not slower than the rules at a load in between, beyond the tie and the noise of the measurements. + +A latent-charge long-range module evaluated by the fused long-range operator +runs its charge network through the same routes, by a table of its own. The +measurement times that network in place of the energy fitting of the fused +operator without a module, which runs the same kernels on runs of the same +sizes; a charge network of the key of the energy fitting shares its table. """ from __future__ import ( @@ -155,7 +161,8 @@ def network_key(fitting: Any) -> tuple | None: Parameters ---------- fitting : Any - The pt_expt energy fitting network. + The pt_expt fitting network: an energy fitting or the charge network + of a long-range module. Returns ------- @@ -196,7 +203,8 @@ def fitting_routes(fitting: Any) -> list[int]: Parameters ---------- fitting : Any - The pt_expt energy fitting network. + The pt_expt fitting network: an energy fitting or the charge network + of a long-range module. Returns ------- @@ -274,7 +282,10 @@ def _bulk_graph(rcut: float, node_count: int, device: torch.device) -> Any: class _Bench: """Times the fused operator of one model on synthetic bulk systems. - The bench keeps four counters, which :func:`tune_fitting_routes` logs: + The operator is that without a long-range module, with the descriptor of + the model and the network ``fitting`` with its atomic bias ``atom_bias``, + the energy fitting or a charge network, in place of the energy fitting. + The bench keeps four counters, which :func:`_tune_network` logs: - ``samples``, the samples of run sizes that it times, those of the bisections of bands included; @@ -287,17 +298,20 @@ class _Bench: those of the trains, their lead-ins included. """ - def __init__(self, model: Any, device: torch.device, largest: int) -> None: + def __init__( + self, + model: Any, + device: torch.device, + largest: int, + fitting: Any, + atom_bias: torch.Tensor, + ) -> None: import copy - from deepmd.pt_expt.model.make_model import ( - fused_atom_bias, - ) - learned, pair_potential = model.atomic_model.fused_decomposition() self.descriptor = copy.deepcopy(learned.descriptor).to(device).eval() - self.fitting = copy.deepcopy(learned.fitting_net).to(device).eval() - self.atom_bias = fused_atom_bias(model.atomic_model, learned).to(device) + self.fitting = copy.deepcopy(fitting).to(device).eval() + self.atom_bias = atom_bias.to(device) self.pair_table = ( None if pair_potential is None else pair_potential.pair_table.to(device) ) @@ -662,38 +676,42 @@ def _measure( return [high for high, _ in merged[:-1]], [choice for _, choice in merged] -def tune_fitting_routes( - model: torch.nn.Module, level: int, device: torch.device -) -> None: - """Measure the fitting routes of a compressed DPA4C model on ``device``. +def _measured_networks(model: torch.nn.Module) -> list[tuple[Any, torch.Tensor]]: + """Return the networks whose routes the fused operators of a model take. - A no-op for a model that does not take the fused energy-force operator of - compressed DPA4C, and for a network and device whose table is already - measured in this process. - - Parameters - ---------- - model : torch.nn.Module - The model about to be frozen. - level : int - The ``DP_TRITON_INFER`` level, which the measurement does not read. - device : torch.device - The CUDA device the frozen model is compiled for. + They are the energy fitting with its combined atomic bias and, for a + long-range module that the fused long-range operator evaluates, its + charge network with its per-type bias. """ - del level - from deepmd.pt_expt.kernels.dpa4c.canonical import ( - canonical_model_eligible, - op_available, + from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + fused_long_range_eligible, + ) + from deepmd.pt_expt.model.make_model import ( + fused_atom_bias, ) + + learned, _ = model.atomic_model.fused_decomposition() + networks = [(learned.fitting_net, fused_atom_bias(model.atomic_model, learned))] + if learned.has_long_range() and fused_long_range_eligible(learned.long_range): + charge_net = learned.long_range.charge_net + networks.append((charge_net, charge_net.bias_atom_e[:, 0].detach())) + return networks + + +def _tune_network( + model: torch.nn.Module, fitting: Any, atom_bias: torch.Tensor, device: torch.device +) -> None: + """Measure the routes of one network of ``model`` on ``device``. + + A no-op for a network and device whose table is already measured in this + process, and for a width of which no fused shape runs on the device. + """ from deepmd.pt_expt.kernels.graph_fitting import ( fitting_operator_arguments, node_tile, ) - if not op_available() or not canonical_model_eligible(model): - return - learned, _ = model.atomic_model.fused_decomposition() - key = network_key(learned.fitting_net) + key = network_key(fitting) if key is None: return fingerprint = _fingerprint(device) @@ -716,7 +734,7 @@ def tune_fitting_routes( # of the freeze current; the search launched on another device than that # of its positions would read them out of bounds. with torch.cuda.device(device), torch.no_grad(): - bench = _Bench(model, device, largest) + bench = _Bench(model, device, largest, fitting, atom_bias) network = fitting_operator_arguments(bench.fitting) bounds, choices = _measure(bench, list(network.weights), largest) table = (ROUTE_TABLE_MAGIC, *fingerprint, len(choices), *bounds, *choices) @@ -725,7 +743,7 @@ def tune_fitting_routes( "Measured the fitting routes of hidden width %d on %s in %.1f s " "(%d samples, %d of them checked on trains of one route and %d of " "these overturned, %d replays): %d intervals, bounds %s, choices %s", - int(network.weights[0].shape[1]), + width, torch.cuda.get_device_name(device), time.perf_counter() - start, bench.samples, @@ -738,6 +756,37 @@ def tune_fitting_routes( ) +def tune_fitting_routes( + model: torch.nn.Module, level: int, device: torch.device +) -> None: + """Measure the fitting routes of a compressed DPA4C model on ``device``. + + The routes of the energy fitting and, for a long-range module that the + fused long-range operator evaluates, of its charge network are measured + one network at a time. A no-op for a model that does not take the fused + energy-force operator of compressed DPA4C. + + Parameters + ---------- + model : torch.nn.Module + The model about to be frozen. + level : int + The ``DP_TRITON_INFER`` level, which the measurement does not read. + device : torch.device + The CUDA device the frozen model is compiled for. + """ + del level + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + canonical_model_eligible, + op_available, + ) + + if not op_available() or not canonical_model_eligible(model): + return + for fitting, atom_bias in _measured_networks(model): + _tune_network(model, fitting, atom_bias, device) + + register_autotuner( "dpa4c_fitting_routes", tune_fitting_routes, diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index f5f73547ac..f18d3e33fc 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -268,7 +268,10 @@ def fused_graph() -> CanonicalGraph: force_mag = None if fused_long_range_eligible(learned.long_range): # The fused operator derives the ownership mask from the frame - # counts itself, as the operator without a long-range module. + # counts itself, and the fitting-route tables measured for the + # energy fitting and the charge network become literals of the + # traced program, as for the operator without a long-range + # module. energy, atom_energy, force, virial, atom_virial, mask = ( dpa4c_canonical_compress_long_range_energy_force( descriptor, @@ -280,6 +283,8 @@ def fused_graph() -> CanonicalGraph: communicator, box, atom_bias, + fitting_routes(fitting), + fitting_routes(learned.long_range.charge_net), ) ) else: diff --git a/doc/model/long-range.md b/doc/model/long-range.md index b883b463fa..81e673febf 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -87,13 +87,13 @@ positive throughout training. charge clouds. It approaches $1/r$ beyond a few $1/\alpha$ and stays finite, $2\alpha/\sqrt{\pi}$, at contact. The inverse width $\alpha$ is set by {ref}`alpha `. -- `sog`: $K(r) = \sum_k a_k \exp(-r^2/2\sigma_k^2)$, a sum of Gaussians whose - amplitudes $a_k$ may take either sign. The widths are given by +- `sog`: $K(r) = \sum_m a_m \exp(-r^2/2\sigma_m^2)$, a sum of Gaussians whose + amplitudes $a_m$ may take either sign. The widths are given by {ref}`width ` or default to a geometric sequence, and the amplitudes by {ref}`amplitude `, which explicit widths require. The default amplitudes, - $a_k = \sqrt{2/\pi}\,\ln b/\sigma_k$ for the ratio $b$ of consecutive + $a_m = \sqrt{2/\pi}\,\ln b/\sigma_m$ for the ratio $b$ of consecutive widths, discretize $1/r$: with the default widths the kernel stays finite, $0.46$ Å⁻¹, at contact, and its radial force between fixed charges stays within 0.4% of that of $1/r$ from 6.5 Å to 150 Å. It differs from $1/r$ @@ -117,7 +117,7 @@ a frame given a cell is evaluated over the images of the cell. The sum converges in reciprocal space, where the Fourier transform of the kernel, - `les`: $\hat K(\mathbf{k}) = (4\pi/k^2)\, e^{-k^2/4\alpha^2}$, or -- `sog`: $\hat K(\mathbf{k}) = \sum_k a_k (2\pi)^{3/2} \sigma_k^3\, e^{-\sigma_k^2 k^2/2}$, +- `sog`: $\hat K(\mathbf{k}) = \sum_m a_m (2\pi)^{3/2} \sigma_m^3\, e^{-\sigma_m^2 k^2/2}$, which with the default amplitudes approaches the Coulomb transform $4\pi/k^2$, diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 53f7091e79..b1507d7103 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -1253,8 +1253,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi { private: /** * @brief Check that every model accepts the cell handed to the ensemble. - * Without one, the neighbor-list entries serve a model with a long-range - * module beside models that take the cell of the region. + * An empty cell, which a non-periodic system hands a model with a + * long-range module, is accepted only when every model carries one: the + * neighbor-list entries of other models, such as those of the TensorFlow + * backend, read the cell of the region. **/ template void check_ensemble_box(const std::vector& dbox) const; diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index dba250c575..9646f9794a 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -286,6 +286,7 @@ void DeepPotPTExpt::init(const std::string& model, nnei += v.as_int(); } } + long_range_ = metadata.obj_val.count("topologies") > 0; if (metadata.obj_val.count("lower_input_kind")) { const std::string lower_input_kind = metadata["lower_input_kind"].as_string(); @@ -295,11 +296,12 @@ void DeepPotPTExpt::init(const std::string& model, lower_input_kind == "dpa4c_canonical"; // Whether the force assembly of the compact lower reads the source CSR on // every call, as the archive records it; an archive without the record - // reads it for compressed DPA1 and for a long-range module (below). + // reads it for compressed DPA1 and for a long-range module. canonical_requires_source_csr_ = metadata.obj_val.count("canonical_requires_source_csr") ? metadata["canonical_requires_source_csr"].as_bool() - : lower_input_kind == "dpa1_canonical"; + : lower_input_kind == "dpa1_canonical" || + (lower_input_is_canonical_ && long_range_); } else { lower_input_is_edge_ = false; lower_input_is_graph_ = false; @@ -309,7 +311,6 @@ void DeepPotPTExpt::init(const std::string& model, // A model with a long-range module reads the all-pairs topology of the // frame after its other inputs, and, when the module serves periodic // systems, the mesh of the cell, sized by the spacing of the metadata. - long_range_ = metadata.obj_val.count("topologies") > 0; long_range_periodic_ = false; if (long_range_) { for (const auto& name : metadata["topologies"].as_array()) { @@ -330,9 +331,6 @@ void DeepPotPTExpt::init(const std::string& model, "a model with a long-range module requires the graph or the " "compact canonical lower."); } - if (!metadata.obj_val.count("canonical_requires_source_csr")) { - canonical_requires_source_csr_ = lower_input_is_canonical_; - } } if (lower_input_is_edge_) { std::cerr << "WARNING: This .pt2 uses the deprecated edge_vec lower " diff --git a/source/api_cc/src/longRangePTExpt.h b/source/api_cc/src/longRangePTExpt.h index f955f1f6f9..65c3c21b6e 100644 --- a/source/api_cc/src/longRangePTExpt.h +++ b/source/api_cc/src/longRangePTExpt.h @@ -73,7 +73,9 @@ inline std::vector buildAllPairsTopology( * pair: the reciprocal-space topology describes the interactions. * * Counterpart of Python ``empty_pair_graph`` for the exported lower: two - * masked guard pairs, the lower bound of its pair axis. + * masked guard pairs, the lower bound of its pair axis. Python pads four, + * which keeps the pair axis of a traced sample apart from the axes of length + * two and three; the exported lower accepts any length from two on. */ inline std::vector buildGuardPairTopology( const torch::Device& device, const torch::ScalarType vec_dtype) { diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 58913df34f..94fb813a3f 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -166,7 +166,7 @@ else: # script also writes the compressed archive, whose compact canonical lower # reads the node positions. env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_long_range.py & - PID16=$! + PID17=$! wait $PID11 wait $PID12 wait $PID13 @@ -176,7 +176,7 @@ else: wait $PID15 wait $PID16 fi - wait $PID16 + wait $PID17 fi if [ "${ENABLE_PADDLE:-TRUE}" == "TRUE" ]; then PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir diff --git a/source/install/test_lammps_kokkos.sh b/source/install/test_lammps_kokkos.sh index b5f2506c16..78a6dd98a4 100755 --- a/source/install/test_lammps_kokkos.sh +++ b/source/install/test_lammps_kokkos.sh @@ -106,3 +106,10 @@ PY DPA4SPIN_KOKKOS_REQUIRED=1 python -m pytest -v -s \ "${REPO_ROOT}/source/lmp/tests/test_dpa4spin_kokkos_magnetic_force.py" +# The compressed long-range fixtures come from gen_dpa4c_long_range.py of +# test_cc_local.sh; the MPI runner of the test launches mpirun from PATH, which +# leads with the launcher of the MPI prefix, while pytest runs on the Python of +# the job. +PYTHON=$(command -v python) +PATH=${MPI_ROOT}/bin:${PATH} DPA4C_LONG_RANGE_KOKKOS_REQUIRED=1 "${PYTHON}" -m pytest -v -s \ + "${REPO_ROOT}/source/lmp/tests/test_lammps_dpa4c_long_range_kokkos.py" diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 2ed86f2da2..cc7c6b1744 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -992,6 +992,18 @@ void PairDeepMD::init_style() { "a system periodic in every dimension or in none; use p in " "every dimension, or f, s or m in every dimension."); } + // A non-periodic system reaches an ensemble that holds a long-range module + // without a cell, which only models with such a module accept. The steps + // without a model deviation evaluate the first model alone; the model + // deviation checks every model at the steps that evaluate them all. + if (numb_models > 1 && long_range && periodic_dimensions == 0 && + !deep_pot.has_long_range()) { + error->all(FLERR, + "pair style deepmd: an ensemble that holds a model with a " + "long-range module evaluates a non-periodic system only when " + "every model carries one; use models with long-range modules " + "throughout, or a periodic system."); + } } /* ---------------------------------------------------------------------- diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index 2d24f2c4e9..4fe527b821 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -414,13 +414,15 @@ template const double* PairDeepMDKokkos::model_node_positions() { atomKK->sync(execution_space, X_MASK); auto x = atomKK->k_x.template view(); - // Relative to the origin of the box, as the plain pair style passes them. - // Without virtual atoms the model nodes are the leading atoms, whose - // coordinates are already those positions in the length unit of the model - // for a box at the origin and a conversion factor of one. - const double x0 = domain->boxlo[0]; - const double y0 = domain->boxlo[1]; - const double z0 = domain->boxlo[2]; + // A long-range module reads the positions relative to the origin of the + // box, as the plain pair style passes them, since its per-atom virial in a + // periodic cell depends on the origin; other models read them only through + // quantities that do not. Without virtual atoms the model nodes are the + // leading atoms, whose coordinates are already those positions in the + // length unit of the model at a conversion factor of one. + const double x0 = long_range ? domain->boxlo[0] : 0.0; + const double y0 = long_range ? domain->boxlo[1] : 0.0; + const double z0 = long_range ? domain->boxlo[2] : 0.0; if (!compact_graph.has_null_types && dist_unit_cvt_factor == 1.0 && x0 == 0.0 && y0 == 0.0 && z0 == 0.0) { return x.data(); diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index 5858a0895f..22e3410132 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -83,9 +83,9 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // neighbor list, returning the edge count. Public because it launches // extended device lambdas, which CUDA forbids inside non-public members. int build_edges_device(); - // Positions of the model nodes relative to the origin of the box, in - // model-node order and in the length unit of the model; public for the same - // reason. + // Positions of the model nodes in model-node order and in the length unit + // of the model, relative to the origin of the box for a model with a + // long-range module; public for the same reason. const double* model_node_positions(); protected: diff --git a/source/lmp/tests/lammps_test_utils.py b/source/lmp/tests/lammps_test_utils.py index e3f33f5cb1..ff73605625 100644 --- a/source/lmp/tests/lammps_test_utils.py +++ b/source/lmp/tests/lammps_test_utils.py @@ -12,6 +12,7 @@ ) import os +import signal import subprocess as sp import sys import tempfile @@ -136,6 +137,7 @@ def run_mpi_pair_runner( runner_args: list[str] | None = None, output_columns: tuple[tuple[str, int], ...] = (("forces", 3), ("virials", 9)), capture: bool = False, + timeout: float | None = None, ) -> dict[str, Any]: """Invoke a DPA MPI runner and parse its energy/per-atom output. @@ -144,7 +146,9 @@ def run_mpi_pair_runner( model-specific wrapper retains its own defaults and explanatory docstring. If ``capture`` is true, skip parsing and return the subprocess result as - ``{"returncode": int, "stdout": str, "stderr": str}``. + ``{"returncode": int, "stdout": str, "stderr": str}``. ``timeout`` bounds + the run in seconds, after which the launcher and its ranks are killed and + ``subprocess.TimeoutExpired`` is raised; ``None`` leaves it unbounded. """ with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: output_path = Path(f.name) @@ -167,15 +171,30 @@ def run_mpi_pair_runner( argv.extend(extra_args) if runner_args: argv.extend(runner_args) + # The launcher and the ranks of a bounded run share a process group of + # their own, so that a timeout ends the ranks with it rather than + # leaving them holding their devices. + with sp.Popen( + argv, + stdout=sp.PIPE if capture else None, + stderr=sp.PIPE if capture else None, + text=True, + start_new_session=timeout is not None, + ) as process: + try: + stdout, stderr = process.communicate(timeout=timeout) + except sp.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.communicate() + raise if capture: - proc = sp.run(argv, capture_output=True, text=True) return { - "returncode": proc.returncode, - "stdout": proc.stdout, - "stderr": proc.stderr, + "returncode": process.returncode, + "stdout": stdout, + "stderr": stderr, } - - sp.check_call(argv) + if process.returncode != 0: + raise sp.CalledProcessError(process.returncode, argv) lines = output_path.read_text().strip().splitlines() rows = np.array( [list(map(float, line.split())) for line in lines[1:]], diff --git a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py index a73908d6f3..55d5317943 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py +++ b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py @@ -56,7 +56,7 @@ type=str, default="p p p", help="LAMMPS boundary style. Default 'p p p'; a model with a long-range " - "module evaluates complete non-periodic systems and passes 'f f f'.", + "module takes 'f f f', or 'p p p' when its module serves periodic systems.", ) parser.add_argument( "--pair-coeff", @@ -134,9 +134,24 @@ "to exercise sendlist composition changes from real-atom motion " "rather than only from NULL motion.", ) +parser.add_argument( + "--kokkos", + action="store_true", + help="Run pair_style deepmd/kk on the Kokkos CUDA device of each rank, " + "with a full neighbor list, newton off and host communication, as the " + "device path of LAMMPS + Kokkos runs it without a CUDA-aware MPI. Rank r " + "takes device r modulo the visible devices, as DeepPotPTExpt does.", +) args = parser.parse_args() -lammps = PyLammps() +cmdargs = None +if args.kokkos: + import torch + + devices = min(comm.Get_size(), torch.cuda.device_count()) + cmdargs = ["-k", "on", "g", str(devices), "-sf", "kk", "-pk", "kokkos"] + cmdargs += ["neigh", "full", "newton", "off", "comm", "host", "gpu/aware", "off"] +lammps = PyLammps(cmdargs=cmdargs) # Force the requested domain decomposition. The default "2 1 1" # combined with the simulation box guarantees nswap > 0 on the C++ # side, so DeepPotPTExpt routes to the with-comm AOTI artifact. Pass @@ -194,7 +209,9 @@ else: lammps.velocity(f"nullgroup set {args.null_vx:.6f} 0.0 0.0 units box") -pair_style_line = f"deepmd {args.PB_FILE}" +# The Kokkos style is named explicitly, so that a plugin without it fails +# instead of falling back to the host style. +pair_style_line = f"{'deepmd/kk' if args.kokkos else 'deepmd'} {args.PB_FILE}" if args.pair_extra: pair_style_line += f" {args.pair_extra}" lammps.pair_style(pair_style_line) diff --git a/source/lmp/tests/test_lammps_dpa4c_long_range_kokkos.py b/source/lmp/tests/test_lammps_dpa4c_long_range_kokkos.py new file mode 100644 index 0000000000..9765fdfe2b --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4c_long_range_kokkos.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test the Kokkos pair style with compressed DPA4C models that carry a long-range module. + +``pair_style deepmd/kk`` evaluates a compressed long-range model through the +device-resident compact canonical lower: it gathers the positions of the +model nodes relative to the origin of the box, passes the cell of a periodic +system, and joins the exchange of the long range on every rank. Run +source/install/test_lammps_kokkos.sh after test_cc_local.sh on a CUDA host: +that job supplies the compressed fixtures of +source/tests/infer/gen_dpa4c_long_range.py and a Kokkos LAMMPS library and +plugin, and sets DPA4C_LONG_RANGE_KOKKOS_REQUIRED=1, so that no missing +prerequisite skips the test silently. + +The six atoms of the fixtures, and their box, are translated away from the +origin: the positions relative to the box, and with them every output of the +references, including the per-atom virial, which depends on the origin of a +periodic cell, stay those of the references. +""" + +from __future__ import ( + annotations, +) + +import importlib.util +import os +import shutil +from pathlib import ( + Path, +) + +import constants +import numpy as np +import pytest +from expected_ref import ( + read_expected_ref, +) +from lammps_test_utils import ( + run_mpi_pair_runner, +) +from write_lmp_data import ( + write_lmp_data, +) + +_INFER = Path(__file__).resolve().parents[2] / "tests" / "infer" +_RUNNER = Path(__file__).with_name("run_mpi_pair_deepmd_dpa3_pt2.py") +# The translation of the atoms and of the box. +_ORIGIN = np.array([-4.0, 2.5, 7.0]) +_COORD = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +# atype [0, 1, 1, 0, 1, 1] (O, H, H, O, H, H) -> LAMMPS types [1, 2, 2, 1, 2, 2] +_TYPE = np.array([1, 2, 2, 1, 2, 2]) +# The box of the non-periodic references, and the triclinic cell of the +# periodic ones (gen_dpa4c_long_range.BOX): xlo xhi ylo yhi zlo zhi xy xz yz. +_BOX = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0], dtype=float) +_PERIODIC_BOX = np.array([0, 13, 0, 12.5, 0, 13.5, 0.6, -0.4, 0.8]) +# LAMMPS centroid/stress/atom reports the per-atom stress, the negative of the +# per-atom virial, with the components xx, yy, zz, xy, xz, yz, yx, zx, zy; these +# columns of it give the nine components of the reference virial. +_TO_REFERENCE = [0, 6, 7, 3, 1, 8, 4, 5, 2] + + +@pytest.fixture(scope="module") +def kokkos_job() -> None: + # The ordinary LAMMPS suite also runs on CUDA, but uses a non-Kokkos wheel. + # Only the dedicated matching library/plugin job opts into this regression. + if os.environ.get("DPA4C_LONG_RANGE_KOKKOS_REQUIRED") != "1": + pytest.skip("Enable the dedicated CUDA Kokkos plugin regression explicitly") + for module in ("lammps", "mpi4py"): + if importlib.util.find_spec(module) is None: + pytest.fail(f"Required Python package missing: {module}") + if shutil.which("mpirun") is None: + pytest.fail("mpirun required for the Kokkos regression") + + +def _translated(box: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """The box and the positions of the atoms, translated by ``_ORIGIN``.""" + shifted = box.copy() + shifted[0:6] += np.repeat(_ORIGIN, 2) + return shifted, _COORD + _ORIGIN + + +@pytest.mark.parametrize("nprocs", [1, 2]) +@pytest.mark.parametrize( + ("stem", "boundary", "section", "pair_extra"), + [ + ("deeppot_dpa4c_long_range_compressed", "f f f", "default", ""), + ( + "deeppot_dpa4c_long_range_compressed", + "f f f", + "explicit", + "charge_spin 1.0 2.0", + ), + ("deeppot_dpa4c_long_range_periodic_compressed", "p p p", "periodic", ""), + ("deeppot_dpa4c_long_range_periodic_compressed", "f f f", "cluster", ""), + ], +) +def test_pair_deepmd_kokkos_long_range( + kokkos_job: None, + tmp_path: Path, + stem: str, + boundary: str, + section: str, + pair_extra: str, + nprocs: int, +) -> None: + model = _INFER / f"{stem}.pt2" + if not model.is_file(): + pytest.fail(f"Generate the compressed long-range fixture first: {model}") + reference = read_expected_ref(_INFER / f"{stem}.expected")[section] + box, coord = _translated(_PERIODIC_BOX if "periodic" in stem else _BOX) + data = tmp_path / "six.lmp" + write_lmp_data(box, coord, _TYPE, data) + runner_args = ["--boundary", boundary, "--kokkos"] + if pair_extra: + runner_args += ["--pair-extra", pair_extra] + result = run_mpi_pair_runner( + _RUNNER, + data, + model, + nprocs=nprocs, + processors="1 1 1" if nprocs == 1 else "2 1 1", + runner_args=runner_args, + timeout=600, + ) + assert result["pe"] == pytest.approx( + float(np.sum(reference["expected_e"])), rel=1e-6 + ) + np.testing.assert_allclose( + result["forces"], reference["expected_f"].reshape(6, 3), rtol=0, atol=1e-5 + ) + np.testing.assert_allclose( + result["virials"][:, _TO_REFERENCE] / constants.nktv2p, + -reference["expected_v"].reshape(6, 9), + rtol=0, + atol=1e-5, + ) diff --git a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py index 4829a91a7d..2048750aa4 100644 --- a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py @@ -4,10 +4,10 @@ The latent-charge module sums over every pair of atoms of a complete non-periodic system, and over the images of the cell of a periodic one when it serves periodic systems, so ``pair_style deepmd`` runs a model whose module -serves no periodic system with non-periodic boundaries only, and any -long-range model with boundaries periodic in every dimension or in none. A -compressed model also runs on several ranks, whose atoms it gathers over MPI; -an uncompressed one runs on a single rank. The models and the reference +serves no periodic system with non-periodic boundaries only, and a model whose +module serves periodic systems with boundaries periodic in every dimension or +in none. A compressed model also runs on several ranks, whose atoms it gathers +over MPI; an uncompressed one runs on a single rank. The models and the reference values are generated by source/tests/infer/gen_dpa4c_long_range.py, the compressed models for CUDA only: without a cell, at the default charge state and at ``charge_spin 1.0 2.0``, for a module that serves no periodic system, diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 3a596618b0..4c7e4e29d1 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -319,9 +319,9 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) if(CMAKE_SYSTEM_NAME STREQUAL "Linux") target_link_options(deepmd_op_pt PRIVATE "LINKER:--no-undefined") endif() - # The direct sum of the long range evaluates every pair with the special - # function unit; without denormal handling each of its three calls is a single - # instruction. + # The pair loop of the long-range direct sum calls exp2f and rsqrtf, which + # fast math, without denormal handling, compiles to a single special-function + # instruction each. set_source_files_properties(long_range/direct_sum.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math") endif() diff --git a/source/op/pt/dpa4c/graph_compress.cu b/source/op/pt/dpa4c/graph_compress.cu index 8eef579e48..48a8aaf062 100644 --- a/source/op/pt/dpa4c/graph_compress.cu +++ b/source/op/pt/dpa4c/graph_compress.cu @@ -1496,6 +1496,273 @@ void mark_broken_frames(const torch::Tensor& energy, } } +// A fitting network of the fused energy-force operators -- the energy fitting +// or the charge network of a long-range module: its layer plan, weights, +// head, per-type bias and activation. +struct FittingNetwork { + const FittingLayerPlan& plan; + const std::vector& ws; + const std::vector& bs; + const std::vector& resnets; + const torch::Tensor& w_head; + const torch::Tensor& b_head; + const torch::Tensor& bias_atom_e; + int64_t act; +}; + +// The fitting routes of the runs of ``run`` nodes into which a fused operator +// splits its node axis, and the buffers they need. The full runs and the tail +// run each take the route of their own size; an absent table, like an empty +// one, leaves both to the rules. A route of the fused network needs the +// transposed weights, and its kernel writes the pre-activations of a run and +// the sheet through which the CTAs of a cluster exchange their columns into +// the pre-activation buffer; the per-layer kernels need the activation +// ping-pong, sized for the largest run that takes them, and the +// pre-activations of a run between their forward and their backward. With +// ``deferred``, every run of the per-layer kernels keeps its pre-activations +// at its own place on the node axis, so that its backward may follow the +// forward of other runs and other work of the operator; the buffer then holds +// every node, as the descriptor of the long-range operator does. +class FittingRuns { + public: + FittingRuns(cudaStream_t stream, + const FittingNetwork& fitting, + long run, + long node_count, + std::optional> table, + const torch::TensorOptions& f32, + bool deferred = false) + : fitting_(fitting), run_(run) { + const std::vector routes = + std::move(table).value_or(std::vector{}); + const long tail = node_count % run; + const long full_count = node_count / run; + full_route_ = fitting_route(fitting.plan, fitting.ws, run, routes); + tail_route_ = tail > 0 + ? fitting_route(fitting.plan, fitting.ws, tail, routes) + : full_route_; + const long per_layer_nodes = + std::max(full_count > 0 && full_route_ == kFittingPerLayer ? run : 0L, + tail > 0 && tail_route_ == kFittingPerLayer ? tail : 0L); + if ((full_count > 0 && full_route_ != kFittingPerLayer) || + (tail > 0 && tail_route_ != kFittingPerLayer)) { + transposed_ = + torch::empty({fitting_fused_transposed_floats(fitting.ws)}, f32); + fitting_fused_transpose(stream, fitting.plan, fitting.ws, + transposed_.data_ptr()); + } + if (per_layer_nodes > 0) { + const int slots = fitting.plan.n_layer > 1 ? 2 : 1; + scratch_ = + torch::empty({slots, per_layer_nodes, fitting.plan.width_max}, f32); + slot_[0] = scratch_[0].data_ptr(); + slot_[1] = slots > 1 ? scratch_[1].data_ptr() : slot_[0]; + } + // The kept pre-activations of the run from node ``begin`` start at + // ``begin * saved_width``, a 16-byte boundary, since every hidden width is + // a multiple of four (check_weights of graph_fitting.cu). + kept_ = deferred && per_layer_nodes > 0; + saved_ = torch::empty( + {(kept_ ? node_count : run) * fitting.plan.saved_width()}, f32); + } + + // The outputs ``output`` of the ``count`` nodes of the run from node + // ``begin``, whose descriptor rows are ``x``, and the cotangent ``d_x`` of + // those rows from the output cotangent ``seed``; ``d_x`` may take the place + // of ``x``. + void output_and_cotangent(cudaStream_t stream, + long begin, + long count, + const float* x, + float* d_x, + long width, + const long* node_type, + const double* seed, + double* output) const { + forward(stream, begin, count, x, d_x, width, node_type, seed, output); + backward(stream, begin, count, seed, d_x); + } + + // The outputs of a run as in output_and_cotangent, and on a route of the + // fused network its cotangent as well, which that network forms in the same + // pass; a run of the per-layer kernels saves its pre-activations for + // backward(). + void forward(cudaStream_t stream, + long begin, + long count, + const float* x, + float* d_x, + long width, + const long* node_type, + const double* seed, + double* output) const { + const FittingNetwork& fitting = fitting_; + const int64_t route = route_of(count); + if (route != kFittingPerLayer) { + fitting_energy_gradient_range( + stream, fitting.plan, x, width, node_type, fitting.ws, fitting.bs, + fitting.resnets, fitting.w_head, fitting.b_head, fitting.bias_atom_e, + fitting.act, transposed_.data_ptr(), seed, count, saved(begin), + output, d_x, route); + return; + } + fitting_forward_range(stream, fitting.plan, x, width, node_type, fitting.ws, + fitting.bs, fitting.resnets, fitting.w_head, + fitting.b_head, fitting.bias_atom_e, fitting.act, + count, saved(begin), slot_, output); + } + + // The cotangent ``d_x`` of a run of the per-layer kernels from ``seed``, + // after its forward(); nothing for a run of the fused network, whose + // forward() has formed it. Without ``deferred`` it follows the forward of + // its own run. + void backward(cudaStream_t stream, + long begin, + long count, + const double* seed, + float* d_x) const { + if (route_of(count) != kFittingPerLayer) { + return; + } + const FittingNetwork& fitting = fitting_; + fitting_backward_range(stream, fitting.plan, seed, saved(begin), fitting.ws, + fitting.bs, fitting.resnets, fitting.w_head, + fitting.act, count, slot_[0], + fitting.plan.n_layer > 1 ? slot_[1] : nullptr, d_x); + } + + private: + int64_t route_of(long count) const { + return count == run_ ? full_route_ : tail_route_; + } + + float* saved(long begin) const { + return saved_.data_ptr() + + (kept_ ? begin * fitting_.plan.saved_width() : 0); + } + + const FittingNetwork& fitting_; + long run_; + int64_t full_route_ = 0; + int64_t tail_route_ = 0; + bool kept_ = false; + torch::Tensor transposed_; + torch::Tensor scratch_; + torch::Tensor saved_; + float* slot_[2] = {nullptr, nullptr}; +}; + +// The state of a fused energy-force evaluation that its epilogue reads. +struct EpilogueInputs { + const torch::Tensor& n_node; + const torch::Tensor& frame_row_ptr; + const torch::Tensor& force_fixed; + const torch::Tensor& raw_energy; + const torch::Tensor& seed; + torch::Tensor& node_virial; + EvaluationWords* evaluation_words; +}; + +// The magnetic terms of a native-spin model: the fixed-point totals and the +// on-site gradient that form the magnetic force. +struct SpinTerms { + const torch::Tensor& fixed; + const torch::Tensor& gradient; + torch::Tensor& force_mag; +}; + +// The inputs of the per-atom virial, which the source CSR attributes to the +// source nodes of the edges. +struct SourceVirialTerms { + const torch::Tensor& edge_gradient; + const torch::Tensor& edge_vec; + const torch::Tensor& row_ptr; + const torch::Tensor& order; +}; + +// The terms of a long-range module: node energies (N, 1) float64, forces at +// fixed charges (N, 3) and per-node virials (N, 9) float32. +struct LongRangeTerms { + const torch::Tensor& energy; + const torch::Tensor& force; + const torch::Tensor& atom_virial; +}; + +// The node results and the frame totals of a fused energy-force evaluation +// (energy_force_epilogue_kernel for a single frame, energy_force_nodes_kernel +// and segment reductions for several), and the per-atom virial +// ``atom_virial_output``, which is zero without ``source_virial``. ``spin`` is +// null for a model without spin and ``long_range`` for a model without a +// long-range module. +void finish_energy_force(cudaStream_t stream, + long node_count, + int blocks, + const EpilogueInputs& step, + const SpinTerms* spin, + const SourceVirialTerms* source_virial, + const LongRangeTerms* long_range, + torch::Tensor& energy, + torch::Tensor& atom_energy, + torch::Tensor& force, + torch::Tensor& virial, + torch::Tensor& atom_virial_output) { + const auto* force_words = reinterpret_cast( + step.force_fixed.data_ptr()); + const auto* spin_words = spin != nullptr + ? reinterpret_cast( + spin->fixed.data_ptr()) + : nullptr; + const float* spin_gradient = + spin != nullptr ? spin->gradient.data_ptr() : nullptr; + double* force_mag = + spin != nullptr ? spin->force_mag.data_ptr() : nullptr; + const double* extra_energy = + long_range != nullptr ? long_range->energy.data_ptr() : nullptr; + const float* extra_force = + long_range != nullptr ? long_range->force.data_ptr() : nullptr; + float* atom_virial_zero = + source_virial != nullptr ? nullptr : atom_virial_output.data_ptr(); + if (step.n_node.size(0) == 1) { + auto partial = torch::empty({blocks, kEpilogueTotals}, + energy.options().dtype(torch::kFloat64)); + energy_force_epilogue_kernel<<>>( + node_count, step.n_node.data_ptr(), force_words, spin_words, + spin_gradient, step.raw_energy.data_ptr(), + step.seed.data_ptr(), step.node_virial.data_ptr(), + extra_energy, extra_force, + long_range != nullptr ? long_range->atom_virial.data_ptr() + : nullptr, + force.data_ptr(), force_mag, atom_energy.data_ptr(), + atom_virial_zero, partial.data_ptr(), step.evaluation_words, + energy.data_ptr(), virial.data_ptr()); + DPA4C_CHECK_LAUNCH("dpa4c energy-force epilogue"); + } else { + energy_force_nodes_kernel<<>>( + node_count, force_words, spin_words, spin_gradient, + step.raw_energy.data_ptr(), step.seed.data_ptr(), + extra_energy, extra_force, force.data_ptr(), force_mag, + atom_energy.data_ptr(), atom_virial_zero); + DPA4C_CHECK_LAUNCH("dpa4c energy-force nodes"); + energy.copy_(frame_scalar_sum(atom_energy, step.n_node)); + if (long_range != nullptr) { + step.node_virial.add_(long_range->atom_virial); + } + virial.copy_(frame_virial_sum(step.node_virial, step.n_node)); + mark_broken_frames(energy, virial, step.frame_row_ptr, node_count, force, + spin != nullptr ? &spin->force_mag : nullptr); + } + if (source_virial != nullptr) { + atom_virial_output.copy_(canonical_source_virial( + source_virial->edge_gradient, source_virial->edge_vec, + source_virial->row_ptr, source_virial->order, node_count)); + if (long_range != nullptr) { + atom_virial_output.add_( + long_range->atom_virial.reshape({node_count, 3, 3})); + } + } +} + } // namespace // Energy, force and virial of one compressed inference step, evaluated over @@ -1671,36 +1938,10 @@ dpa4c_canonical_compress_energy_force( : node_count; auto descriptor = torch::empty({run, widths.output_width}, f32); auto state = torch::empty({run, widths.moment_width + 2}, f32); - auto saved = torch::empty({run * plan.saved_width()}, f32); - // The full runs and the tail run each take the fitting route of their own - // size; an absent table, like an empty one, leaves both to the rules. A - // route of the fused network needs only the transposed weights besides the - // saved pre-activations; the per-layer kernels need the activation - // ping-pong, sized for the largest run that takes them. - const std::vector routes = - std::move(fitting_routes).value_or(std::vector{}); - const long tail = node_count % run; - const int64_t full_route = fitting_route(plan, ws, run, routes); - const int64_t tail_route = - tail > 0 ? fitting_route(plan, ws, tail, routes) : full_route; - const long full_count = node_count / run; - const long per_layer_nodes = - std::max(full_count > 0 && full_route == kFittingPerLayer ? run : 0L, - tail > 0 && tail_route == kFittingPerLayer ? tail : 0L); - torch::Tensor transposed; - torch::Tensor scratch; - float* slot[2] = {nullptr, nullptr}; - if ((full_count > 0 && full_route != kFittingPerLayer) || - (tail > 0 && tail_route != kFittingPerLayer)) { - transposed = torch::empty({fitting_fused_transposed_floats(ws)}, f32); - fitting_fused_transpose(stream, plan, ws, transposed.data_ptr()); - } - if (per_layer_nodes > 0) { - const int slots = plan.n_layer > 1 ? 2 : 1; - scratch = torch::empty({slots, per_layer_nodes, plan.width_max}, f32); - slot[0] = scratch[0].data_ptr(); - slot[1] = slots > 1 ? scratch[1].data_ptr() : slot[0]; - } + const FittingNetwork fitting{plan, ws, bs, resnets, + w_head, b_head, bias_atom_e, act}; + const FittingRuns runs(stream, fitting, run, node_count, + std::move(fitting_routes), f32); auto empty_index = torch::empty({0}, source.options()); auto empty_mask = torch::empty({0}, edge_vec.options().dtype(torch::kBool)); @@ -1750,26 +1991,11 @@ dpa4c_canonical_compress_energy_force( dispatch(channels, false, arguments, stream); // The cotangent replaces the descriptor, which the run no longer needs. - const int64_t route = count == run ? full_route : tail_route; - if (route != kFittingPerLayer) { - fitting_energy_gradient_range( - stream, plan, descriptor.data_ptr(), widths.output_width, - node_type.data_ptr() + begin, ws, bs, resnets, w_head, b_head, - bias_atom_e, act, transposed.data_ptr(), - seed.data_ptr() + begin, count, saved.data_ptr(), - raw_energy.data_ptr() + begin, descriptor.data_ptr(), - route); - } else { - fitting_forward_range( - stream, plan, descriptor.data_ptr(), widths.output_width, - node_type.data_ptr() + begin, ws, bs, resnets, w_head, b_head, - bias_atom_e, act, count, saved.data_ptr(), slot, - raw_energy.data_ptr() + begin); - fitting_backward_range( - stream, plan, seed.data_ptr() + begin, - saved.data_ptr(), ws, bs, resnets, w_head, act, count, slot[0], - plan.n_layer > 1 ? slot[1] : nullptr, descriptor.data_ptr()); - } + runs.output_and_cotangent( + stream, begin, count, descriptor.data_ptr(), + descriptor.data_ptr(), widths.output_width, + node_type.data_ptr() + begin, seed.data_ptr() + begin, + raw_energy.data_ptr() + begin); arguments.descriptor_gradient = descriptor.data_ptr(); arguments.state = state.data_ptr(); @@ -1798,47 +2024,15 @@ dpa4c_canonical_compress_energy_force( dispatch(channels, true, arguments, stream); } - const auto* force_words = reinterpret_cast( - force_fixed.data_ptr()); - const auto* spin_words = has_spin - ? reinterpret_cast( - spin_fixed.data_ptr()) - : nullptr; - float* atom_virial_zero = - atom_virial ? nullptr : atom_virial_output.data_ptr(); - if (frame_count == 1) { - auto partial = torch::empty({blocks, kEpilogueTotals}, f64); - energy_force_epilogue_kernel<<>>( - node_count, n_node.data_ptr(), force_words, spin_words, - has_spin ? spin_gradient.data_ptr() : nullptr, - raw_energy.data_ptr(), seed.data_ptr(), - node_virial.data_ptr(), nullptr, nullptr, nullptr, - force.data_ptr(), - has_spin ? force_mag.data_ptr() : nullptr, - atom_energy.data_ptr(), atom_virial_zero, - partial.data_ptr(), evaluation_words, energy.data_ptr(), - virial.data_ptr()); - DPA4C_CHECK_LAUNCH("dpa4c energy-force epilogue"); - } else { - energy_force_nodes_kernel<<>>( - node_count, force_words, spin_words, - has_spin ? spin_gradient.data_ptr() : nullptr, - raw_energy.data_ptr(), seed.data_ptr(), nullptr, - nullptr, force.data_ptr(), - has_spin ? force_mag.data_ptr() : nullptr, - atom_energy.data_ptr(), atom_virial_zero); - DPA4C_CHECK_LAUNCH("dpa4c energy-force nodes"); - energy.copy_(frame_scalar_sum(atom_energy, n_node)); - virial.copy_(frame_virial_sum(node_virial, n_node)); - mark_broken_frames(energy, virial, frame_row_ptr, node_count, force, - has_spin ? &force_mag : nullptr); - } - if (atom_virial) { - atom_virial_output.copy_( - canonical_source_virial(edge_gradient, edge_vec_float, source_row_ptr, - source_order, node_count)); - } + const EpilogueInputs step{n_node, frame_row_ptr, force_fixed, raw_energy, + seed, node_virial, evaluation_words}; + const SpinTerms spin_terms{spin_fixed, spin_gradient, force_mag}; + const SourceVirialTerms source_terms{edge_gradient, edge_vec_float, + source_row_ptr, source_order}; + finish_energy_force(stream, node_count, blocks, step, + has_spin ? &spin_terms : nullptr, + atom_virial ? &source_terms : nullptr, nullptr, energy, + atom_energy, force, virial, atom_virial_output); return {energy, atom_energy, force, virial, atom_virial_output, force_mag, mask}; } @@ -1849,19 +2043,26 @@ dpa4c_canonical_compress_energy_force( // The module reads charges that its network predicts from the descriptor of // every node, and the energy of the charges depends on all of them, so the // descriptor backward cannot start before every charge is known. The -// operator evaluates the descriptor of every node with its moment state, the -// charges from the descriptor, and the long range of the charges at fixed -// charges with its derivative with respect to them -// (deepmd::latent_charge_sum); then the energy fitting with its cotangent and -// the cotangent of the charge network, seeded with that derivative, which it -// adds; and one descriptor backward from the sum of the two cotangents. The -// two networks evaluate the node axis in runs of ``tile`` nodes, as in -// dpa4c_canonical_compress_energy_force, which bounds the memory of their -// activations by the run, the pre-activations of the charge network, which -// its backward reads, by the node axis. The edge backward adds every force -// onto the fixed-point totals as in dpa4c_canonical_compress_energy_force, -// and the epilogue adds the energies, the forces at fixed charges and the -// per-node virials of the long range. +// operator evaluates the descriptor of every node with its moment state, and +// the charges from the descriptor; then the long range of the charges at +// fixed charges with its derivative with respect to them +// (deepmd::latent_charge_sum); then the energy fitting with its cotangent, to +// which the Jacobian row of every charge with respect to the descriptor of +// its node, the cotangent of a unit seed, scaled by that derivative adds the +// cotangent of the charge network; and one descriptor backward from the sum. +// As in dpa4c_canonical_compress_energy_force, the two networks evaluate the +// node axis in runs of ``tile`` nodes, and each takes the fused network or +// the per-layer kernels by a fitting-route table of its own +// (``fitting_routes`` and ``charge_fitting_routes``). A run of the fused +// charge network forms its Jacobian rows in the pass of its charges, and a +// run of the per-layer kernels after the energy fitting, from the +// pre-activations that it keeps for every node; on a V100 at 512 atoms that +// order shortens the LAMMPS step by 3.3% against forming them before the +// long range. The Jacobian rows take the width of the descriptor for every +// node. The edge backward adds every force onto the +// fixed-point totals as in dpa4c_canonical_compress_energy_force, and the +// epilogue adds the energies, the forces at fixed charges and the per-node +// virials of the long range. // // The charge network has one channel and the layout of an eligible energy // fitting; its per-type bias is ``charge_bias``. ``coord``, ``total_charge``, @@ -1935,7 +2136,8 @@ dpa4c_canonical_compress_long_range_energy_force( double prefactor, int64_t order, double spacing, - std::optional> fitting_routes) { + std::optional> fitting_routes, + std::optional> charge_fitting_routes) { constexpr const char* kOp = "dpa4c_canonical_compress_long_range_energy_force"; const bool atom_virial = check_energy_force_inputs( @@ -2065,34 +2267,35 @@ dpa4c_canonical_compress_long_range_energy_force( present_types(arguments, &evaluation_words->node_types); dispatch(channels, false, arguments, stream); - // === Step 2. Charges of the nodes, run by run === - // The charge network reads the descriptor before the energy fitting - // overwrites it with its cotangent. The pre-activations of every run stay, - // run after run in the layout of their run, for the backward of the run. + // === Step 2. Charges of the nodes and their descriptor Jacobian === + // The charge of a node is a function of the descriptor of that node alone, + // so the cotangent of its descriptor through the charge network is the + // derivative of the energy with respect to the charge times the Jacobian + // row of the charge, the input cotangent of a unit seed. A run of the fused + // network forms the charges and those rows in one pass, before the long + // range forms the derivative and the energy fitting overwrites the + // descriptor; a run of the per-layer kernels forms its charges here and + // its rows from their saved pre-activations after the energy fitting. const long run = tile > 0 ? std::max(1, std::min(tile, node_count)) : node_count; const long width = widths.output_width; - const long charge_saved_width = charge_plan.saved_width(); - auto charge_saved = torch::empty({node_count * charge_saved_width}, f32); - const int charge_slots = charge_plan.n_layer > 1 ? 2 : 1; - auto charge_scratch = - torch::empty({charge_slots, run, charge_plan.width_max}, f32); - float* const charge_slot[2] = {charge_scratch[0].data_ptr(), - charge_slots > 1 - ? charge_scratch[1].data_ptr() - : charge_scratch[0].data_ptr()}; - auto raw_charge = torch::empty({node_count}, f64); const auto charge_bias_double = charge_bias.to(torch::kFloat64).contiguous(); + const FittingNetwork charge_fitting{ + charge_plan, charge_ws, charge_bs, charge_resnets, + charge_w_head, charge_b_head, charge_bias_double, charge_act}; + const FittingRuns charge_runs(stream, charge_fitting, run, node_count, + std::move(charge_fitting_routes), f32, true); + const auto unit_seed = torch::ones({run}, f64); + auto raw_charge = torch::empty({node_count}, f64); + auto charge_jacobian = torch::empty({node_count, width}, f32); for (long begin = 0; begin < node_count; begin += run) { const long count = std::min(run, node_count - begin); - fitting_forward_range( - stream, charge_plan, descriptor.data_ptr() + begin * width, - width, node_type.data_ptr() + begin, charge_ws, charge_bs, - charge_resnets, charge_w_head, charge_b_head, charge_bias_double, - charge_act, count, - charge_saved.data_ptr() + begin * charge_saved_width, - charge_slot, raw_charge.data_ptr() + begin); + charge_runs.forward( + stream, begin, count, descriptor.data_ptr() + begin * width, + charge_jacobian.data_ptr() + begin * width, width, + node_type.data_ptr() + begin, unit_seed.data_ptr(), + raw_charge.data_ptr() + begin); } // === Step 3. Long range of the charges === @@ -2107,70 +2310,29 @@ dpa4c_canonical_compress_long_range_energy_force( const torch::Tensor& long_range_force = std::get<2>(long_range); const torch::Tensor& long_range_atom_virial = std::get<3>(long_range); - // === Step 4. Cotangents of the two networks, run by run === - // The full runs and the tail run each take the fitting route of their own - // size, as in dpa4c_canonical_compress_energy_force. - const std::vector routes = - std::move(fitting_routes).value_or(std::vector{}); - const long tail = node_count % run; - const long full_count = node_count / run; - const int64_t full_route = fitting_route(plan, ws, run, routes); - const int64_t tail_route = - tail > 0 ? fitting_route(plan, ws, tail, routes) : full_route; - const long per_layer_nodes = - std::max(full_count > 0 && full_route == kFittingPerLayer ? run : 0L, - tail > 0 && tail_route == kFittingPerLayer ? tail : 0L); - auto saved = torch::empty({run * plan.saved_width()}, f32); - torch::Tensor transposed; - torch::Tensor scratch; - float* slot[2] = {nullptr, nullptr}; - if ((full_count > 0 && full_route != kFittingPerLayer) || - (tail > 0 && tail_route != kFittingPerLayer)) { - transposed = torch::empty({fitting_fused_transposed_floats(ws)}, f32); - fitting_fused_transpose(stream, plan, ws, transposed.data_ptr()); - } - if (per_layer_nodes > 0) { - const int slots = plan.n_layer > 1 ? 2 : 1; - scratch = torch::empty({slots, per_layer_nodes, plan.width_max}, f32); - slot[0] = scratch[0].data_ptr(); - slot[1] = slots > 1 ? scratch[1].data_ptr() : slot[0]; - } - const auto charge_seed = - charge_gradient.select(1, 0).to(torch::kFloat64).contiguous(); - auto charge_cotangent = torch::empty({run, width}, f32); + // === Step 4. Cotangent of the descriptor === + // The cotangent of the energy fitting replaces the descriptor run by run; + // the runs of the per-layer kernels of the charge network then form their + // Jacobian rows, and the rows, scaled by the derivative of the energy with + // respect to the charges, add the cotangent of the charge network. + const FittingNetwork fitting{plan, ws, bs, resnets, + w_head, b_head, bias_atom_e, act}; + const FittingRuns runs(stream, fitting, run, node_count, + std::move(fitting_routes), f32); for (long begin = 0; begin < node_count; begin += run) { const long count = std::min(run, node_count - begin); - // The cotangent of the energy fitting replaces the descriptor of the run, - // which the charge network has read. float* const run_descriptor = descriptor.data_ptr() + begin * width; - const int64_t route = count == run ? full_route : tail_route; - if (route != kFittingPerLayer) { - fitting_energy_gradient_range( - stream, plan, run_descriptor, width, - node_type.data_ptr() + begin, ws, bs, resnets, w_head, b_head, - bias_atom_e, act, transposed.data_ptr(), - seed.data_ptr() + begin, count, saved.data_ptr(), - raw_energy.data_ptr() + begin, run_descriptor, route); - } else { - fitting_forward_range(stream, plan, run_descriptor, width, - node_type.data_ptr() + begin, ws, bs, resnets, - w_head, b_head, bias_atom_e, act, count, - saved.data_ptr(), slot, - raw_energy.data_ptr() + begin); - fitting_backward_range( - stream, plan, seed.data_ptr() + begin, - saved.data_ptr(), ws, bs, resnets, w_head, act, count, slot[0], - plan.n_layer > 1 ? slot[1] : nullptr, run_descriptor); - } - fitting_backward_range( - stream, charge_plan, charge_seed.data_ptr() + begin, - charge_saved.data_ptr() + begin * charge_saved_width, charge_ws, - charge_bs, charge_resnets, charge_w_head, charge_act, count, - charge_slot[0], charge_plan.n_layer > 1 ? charge_slot[1] : nullptr, - charge_cotangent.data_ptr()); - descriptor.slice(0, begin, begin + count) - .add_(charge_cotangent.slice(0, 0, count)); + runs.output_and_cotangent( + stream, begin, count, run_descriptor, run_descriptor, width, + node_type.data_ptr() + begin, seed.data_ptr() + begin, + raw_energy.data_ptr() + begin); + } + for (long begin = 0; begin < node_count; begin += run) { + const long count = std::min(run, node_count - begin); + charge_runs.backward(stream, begin, count, unit_seed.data_ptr(), + charge_jacobian.data_ptr() + begin * width); } + descriptor.addcmul_(charge_jacobian, charge_gradient); // === Step 5. Descriptor backward and force assembly === arguments.descriptor_gradient = descriptor.data_ptr(); @@ -2185,43 +2347,15 @@ dpa4c_canonical_compress_long_range_energy_force( dispatch(channels, true, arguments, stream); // === Step 6. Node results and frame totals === - const auto* force_words = reinterpret_cast( - force_fixed.data_ptr()); - float* atom_virial_zero = - atom_virial ? nullptr : atom_virial_output.data_ptr(); - if (frame_count == 1) { - auto partial = torch::empty({blocks, kEpilogueTotals}, f64); - energy_force_epilogue_kernel<<>>( - node_count, n_node.data_ptr(), force_words, nullptr, nullptr, - raw_energy.data_ptr(), seed.data_ptr(), - node_virial.data_ptr(), long_range_energy.data_ptr(), - long_range_force.data_ptr(), - long_range_atom_virial.data_ptr(), force.data_ptr(), - nullptr, atom_energy.data_ptr(), atom_virial_zero, - partial.data_ptr(), evaluation_words, energy.data_ptr(), - virial.data_ptr()); - DPA4C_CHECK_LAUNCH("dpa4c long-range energy-force epilogue"); - } else { - energy_force_nodes_kernel<<>>( - node_count, force_words, nullptr, nullptr, - raw_energy.data_ptr(), seed.data_ptr(), - long_range_energy.data_ptr(), - long_range_force.data_ptr(), force.data_ptr(), nullptr, - atom_energy.data_ptr(), atom_virial_zero); - DPA4C_CHECK_LAUNCH("dpa4c long-range energy-force nodes"); - energy.copy_(frame_scalar_sum(atom_energy, n_node)); - node_virial.add_(long_range_atom_virial); - virial.copy_(frame_virial_sum(node_virial, n_node)); - mark_broken_frames(energy, virial, frame_row_ptr, node_count, force, - nullptr); - } - if (atom_virial) { - atom_virial_output.copy_( - canonical_source_virial(edge_gradient, edge_vec_float, source_row_ptr, - source_order, node_count)); - atom_virial_output.add_(long_range_atom_virial.reshape({node_count, 3, 3})); - } + const EpilogueInputs step{n_node, frame_row_ptr, force_fixed, raw_energy, + seed, node_virial, evaluation_words}; + const SourceVirialTerms source_terms{edge_gradient, edge_vec_float, + source_row_ptr, source_order}; + const LongRangeTerms long_range_terms{long_range_energy, long_range_force, + long_range_atom_virial}; + finish_energy_force(stream, node_count, blocks, step, nullptr, + atom_virial ? &source_terms : nullptr, &long_range_terms, + energy, atom_energy, force, virial, atom_virial_output); return {energy, atom_energy, force, virial, atom_virial_output, mask}; } diff --git a/source/op/pt/dpa4c/ops.cc b/source/op/pt/dpa4c/ops.cc index 2bd4a0c04e..c249513fbe 100644 --- a/source/op/pt/dpa4c/ops.cc +++ b/source/op/pt/dpa4c/ops.cc @@ -155,7 +155,7 @@ TORCH_LIBRARY_FRAGMENT(deepmd, library) { "int charge_act, Tensor coord, Tensor total_charge, Tensor? box, " "Tensor coefficient, Tensor communicator, int kernel, int channel, " "float prefactor, int order, float spacing, " - "int[]? fitting_routes=None) " + "int[]? fitting_routes=None, int[]? charge_fitting_routes=None) " "-> (Tensor energy, Tensor atom_energy, Tensor force, Tensor virial, " "Tensor atom_virial, Tensor mask)"); } diff --git a/source/op/pt/graph_fitting.cu b/source/op/pt/graph_fitting.cu index 6b1e147533..dec91756da 100644 --- a/source/op/pt/graph_fitting.cu +++ b/source/op/pt/graph_fitting.cu @@ -2402,6 +2402,13 @@ bool launch_first_listed(long panels, Fn&& launch, ShapeList) { template struct FusedShapes; template <> +struct FusedShapes<64> { + using Latency = ShapeList, + FusedShape<64, 1, 4, 4, 2, 2, 1, 8>>; + using Throughput = ShapeList, + FusedShape<64, 1, 4, 4, 4, 2, 2, 2, 16>>; +}; +template <> struct FusedShapes<96> { using Latency = ShapeList, FusedShape<96, 1, 4, 4, 2, 3, 1, 8>>; @@ -2441,6 +2448,8 @@ struct FusedShapes<384> { template bool with_fused_shapes(long width, Fn&& fn) { switch (width) { + case 64: + return fn(FusedShapes<64>{}); case 96: return fn(FusedShapes<96>{}); case 192: @@ -2499,21 +2508,21 @@ bool visit_compared_shapes(long width, bool latency, Fn&& visit) { /// /// A run of more panels than multiprocessors takes the throughput shape of its /// width, the fastest on sm_120 of those that fit the register and -/// shared-memory budgets: the two narrow widths take 32-node panels, three -/// and two CTAs per SM, whose small lane tiles keep the most panels in -/// flight, and a width of 256 takes 64-node panels on one CTA per SM, which -/// stream each weight slice to half as many CTAs as 32-node panels and so keep -/// the L2 traffic within the FMA rate. A width of 96 moves on to 64-node -/// panels of twelve warps, two CTAs per SM in place of three, once a run -/// fills two waves of them, or once they pass it in fewer waves than the -/// 32-node panels, each shape at the CTAs per SM that the device keeps -/// resident. Such a panel streams each weight slice once for twice the nodes -/// of a 32-node panel, and on sm_120 the larger panels take 6% less time at -/// 100,800 nodes and 11% less at 131,072, the default node tile; from two -/// waves on, this gain outweighs the cost of their last, partly filled wave, -/// and below two waves a wave fewer outweighs it: on the 92 multiprocessors -/// of an L20, one wave of 64-node panels takes 0.064 ms at 10,296 nodes -/// against 0.078 ms for two waves of 32-node panels. +/// shared-memory budgets: the narrow widths take 32-node panels, four CTAs +/// per SM at a width of 64, three at 96 and two at 192, whose small lane +/// tiles keep the most panels in flight, and a width of 256 takes 64-node +/// panels on one CTA per SM, which stream each weight slice to half as many +/// CTAs as 32-node panels and so keep the L2 traffic within the FMA rate. A +/// width of 96 moves on to 64-node panels of twelve warps, two CTAs per SM in +/// place of three, once a run fills two waves of them, or once they pass it in +/// fewer waves than the 32-node panels, each shape at the CTAs per SM that the +/// device keeps resident. Such a panel streams each weight slice once for twice +/// the nodes of a 32-node panel, and on sm_120 the larger panels take 6% less +/// time at 100,800 nodes and 11% less at 131,072, the default node tile; from +/// two waves on, this gain outweighs the cost of their last, partly filled +/// wave, and below two waves a wave fewer outweighs it: on the 92 +/// multiprocessors of an L20, one wave of 64-node panels takes 0.064 ms at +/// 10,296 nodes against 0.078 ms for two waves of 32-node panels. /// /// A width that the 128 columns of the per-layer large tile divide is left to /// the per-layer kernels at runs of more panels than multiprocessors on a diff --git a/source/op/pt/long_range/charge_sum.h b/source/op/pt/long_range/charge_sum.h index 4d439a250c..dcccb69393 100644 --- a/source/op/pt/long_range/charge_sum.h +++ b/source/op/pt/long_range/charge_sum.h @@ -22,6 +22,14 @@ namespace deepmd { namespace long_range { +// The channel count and the Gaussian count that the kernels are instantiated +// for; the Gaussians of a kernel are staged in shared memory. +constexpr int kMaxChannels = 8; +constexpr int kMaxGaussians = 32; + +// The interaction kernel: erf(alpha r) / r (LES) or a sum of Gaussians (SOG). +enum KernelKind : int { kLes = 0, kSog = 1 }; + // The atomic energy (owners, 1) float64, the charge gradient (owners, C), the // force (owners, 3), the atomic virial (owners, 9) and the virial of the // owners of every frame (F, 9), float32. @@ -31,8 +39,8 @@ using ChargeSumOutputs = std::tuple; -// The direct sum. ``coefficient`` holds the float64 coefficients of the -// kernel on the device. +// The direct sum, of at least one node and one frame. ``coefficient`` holds +// the float64 coefficients of the kernel on the device. ChargeSumOutputs direct_sum(const torch::Tensor& coord, const torch::Tensor& charge, const torch::Tensor& node_mask, @@ -46,7 +54,8 @@ ChargeSumOutputs direct_sum(const torch::Tensor& coord, int owner_count, const torch::Tensor* communicator); -// The mesh sum. ``box`` holds the cell of every frame, float64 on the host; +// The mesh sum, of at least one node and one frame. ``box`` holds the cell of +// every frame, float64 on the host; // ``coefficient`` the float64 coefficients of the kernel on the device. ChargeSumOutputs mesh_sum(const torch::Tensor& coord, const torch::Tensor& charge, diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu index cb7c56ae93..677f2e25d7 100644 --- a/source/op/pt/long_range/direct_sum.cu +++ b/source/op/pt/long_range/direct_sum.cu @@ -18,9 +18,9 @@ // the difference of the tails, keeps the float32 precision of the // displacement itself wherever the frame lies, while every pair evaluates in // float32; the far field of the erf kernel forms it from the heads alone (see -// evaluate_kernel). For the erf kernel the positions are scaled by alpha, so -// that the pair loop evaluates erf(x) / x and its slope without a -// coefficient, and the reduction restores the units. +// far_batch). For the erf kernel the positions are scaled by alpha, so that +// the pair loop evaluates erf(x) / x and its slope without a coefficient, and +// the reduction restores the units. // // A thread sums one tile in float32 and adds the tile sums with compensated // float32 arithmetic. The partner axis is split into contiguous chunks, so @@ -52,6 +52,8 @@ #include #include #include +#include +#include #include #include "charge_sum.h" @@ -69,7 +71,6 @@ constexpr int kWideTile = 128; constexpr int kNarrowTile = 32; // Partners evaluated together, which share every coefficient load. constexpr int kBatch = 4; -constexpr int kMaxGaussians = 32; // Fills of the resident block slots of the device that one launch aims for. constexpr int kWaves = 8; // Tiles per chunk of wide tiles where a frame that spans every partner takes @@ -83,7 +84,10 @@ constexpr int kWideChunkTiles = 2; // on its own would. constexpr int kNarrowChunkTiles = 2; -enum KernelKind : int { kLes = 0, kSog = 1 }; +using deepmd::long_range::kLes; +using deepmd::long_range::kMaxChannels; +using deepmd::long_range::kMaxGaussians; +using deepmd::long_range::kSog; // Sums of one node: the potential of every channel, the force direction and // the six independent components of the virial. @@ -188,10 +192,12 @@ __global__ void prepare_nodes(const double* __restrict__ coord, // do, and forms the displacements of the batch from the heads alone. The // omitted tails move a displacement by at most 2^-24 (|x_i| + |x_j|), with x // the alpha-scaled offset of a node from the first node of its frame: the -// rounding of float32 positions over the extent of the frame, against a -// distance of at least sqrt(kFarSquared). For a cluster of 100,000 carbon -// atoms at alpha = 0.5 per angstrom, with offsets of up to 26 units, the -// potential and the force keep the error of the exact displacements. +// rounding of float32 positions over the extent of the frame. Against a +// distance of at least sqrt(kFarSquared) = 4.5, the relative error of the +// displacement stays below 2^-24 (|x_i| + |x_j|) / 4.5. For a cluster of +// 100,000 carbon atoms at alpha = 0.5 per angstrom, with offsets of up to 26 +// units, the potential and the force keep the error of the exact +// displacements. // // The sum of Gaussians evaluates sum_k a_k 2^(-c_k r^2) and its slope // sum_k s_k 2^(-c_k r^2), s_k = -2 ln(2) c_k a_k, from the staged @@ -761,18 +767,23 @@ PairLaunch pair_launch(int owner_count, kNarrowChunkTiles}; } +// The blocks of wide tiles that a multiprocessor of the current device keeps +// resident. The occupancy is a property of the compiled kernel and of the +// device, so each device of the process is queried once. template int resident_blocks_per_multiprocessor() { - // The occupancy is a property of the compiled kernel and of the device; a - // process drives one device architecture. - static const int blocks = [] { - int value = 0; - C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &value, direct_sum_kernel, kWideTile, - 0)); - return std::max(1, value); - }(); - return blocks; + static std::mutex mutex; + static std::map blocks; + const int device = c10::cuda::current_device(); + const std::lock_guard lock(mutex); + const auto known = blocks.find(device); + if (known != blocks.end()) { + return known->second; + } + int value = 0; + C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &value, direct_sum_kernel, kWideTile, 0)); + return blocks.emplace(device, std::max(1, value)).first->second; } template @@ -821,6 +832,9 @@ void launch(const DirectSumBuffers& buffers, } } +static_assert(kMaxChannels == 8, + "dispatch_channels instantiates the channel counts 1 to 8"); + template void dispatch_channels(int channels, const DirectSumBuffers& buffers, @@ -867,8 +881,8 @@ namespace long_range { // Evaluates the owner nodes [owner_begin, owner_begin + owner_count) against // the nodes of their frames. Returns the outputs of the owners, row r for // owner owner_begin + r, and the per-frame virial of the owners. Owners other -// than every node lie in a single frame. A non-null ``communicator`` spans -// the domains whose owners share that frame, over which the potential of the +// than every node lie in a single frame. A non-null ``communicator`` spans the +// domains whose owners share that frame, over which the potential of the // constraint is summed. ChargeSumOutputs direct_sum(const torch::Tensor& coord, const torch::Tensor& charge, @@ -889,17 +903,7 @@ ChargeSumOutputs direct_sum(const torch::Tensor& coord, const int64_t channels = charge.size(1); const auto float_options = charge.options(); const auto double_options = coord.options(); - if (node_count == 0 || frame_count == 0) { - return {torch::zeros({owner_count, 1}, double_options), - torch::zeros({owner_count, channels}, float_options), - torch::zeros({owner_count, 3}, float_options), - torch::zeros({owner_count, 9}, float_options), - torch::zeros({frame_count, 9}, float_options)}; - } const bool every_node = owner_begin == 0 && owner_count == node_count; - TORCH_CHECK(every_node || frame_count == 1, - "latent_charge_sum: owners other than every node lie in " - "one frame"); const bool constrained = channel >= 0; const int constrained_channel = static_cast(channel); // The kernels write every entry of the outputs. @@ -910,14 +914,8 @@ ChargeSumOutputs direct_sum(const torch::Tensor& coord, auto virial = torch::empty({frame_count, 9}, float_options); const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); - const auto slices_for = [frame_count](int64_t nodes) { - const int64_t average_nodes = (nodes + frame_count - 1) / frame_count; - return static_cast(std::min( - kMaxSlices, - std::max(1, (average_nodes + kSliceNodes - 1) / kSliceNodes))); - }; - const int slice_count = slices_for(node_count); - const int owner_slice_count = slices_for(owner_count); + const int slice_count = frame_slices(node_count, frame_count); + const int owner_slice_count = frame_slices(owner_count, frame_count); const int components = constrained ? 10 : 9; Workspace workspace; const int64_t frame_ptr_at = workspace.reserve(frame_count + 1); @@ -992,17 +990,10 @@ ChargeSumOutputs direct_sum(const torch::Tensor& coord, owner_count, prefactor, stream.stream()); } } - // The frames of the owners, on the axis of the owners: the frames - // themselves when every node is an owner, else the one frame the owners - // share, written on the device. - torch::Tensor owner_frame_tensor; - const int64_t* owner_frame_ptr = frame_ptr; - if (!every_node) { - owner_frame_tensor = - torch::full({2}, int64_t{owner_count}, counts.options()); - owner_frame_tensor.select(0, 0).zero_(); - owner_frame_ptr = owner_frame_tensor.data_ptr(); - } + torch::Tensor owner_frame_storage; + const int64_t* owner_frame_ptr = + owner_frame_offsets(every_node, frame_ptr, owner_count, counts.options(), + owner_frame_storage); const bool* owner_mask = node_mask.data_ptr() + owner_begin; const dim3 owner_slice_grid(frame_count, owner_slice_count); if (constrained) { diff --git a/source/op/pt/long_range/domain_exchange.cc b/source/op/pt/long_range/domain_exchange.cc index 2ae3d06e05..cbe4f07a23 100644 --- a/source/op/pt/long_range/domain_exchange.cc +++ b/source/op/pt/long_range/domain_exchange.cc @@ -57,9 +57,9 @@ bool reads_device_memory() { } // namespace -Domain domain_of(const torch::Tensor& communicator) { +int domain_count(const torch::Tensor& communicator) { if (address_of(communicator) == 0) { - return {0, 1}; + return 1; } #ifdef USE_MPI int initialized = 0; @@ -67,11 +67,9 @@ Domain domain_of(const torch::Tensor& communicator) { TORCH_CHECK(initialized != 0, "latent_charge_sum: a communicator was given, but MPI is " "not initialized"); - const MPI_Comm comm = communicator_of(communicator); - Domain domain{0, 1}; - MPI_Comm_rank(comm, &domain.rank); - MPI_Comm_size(comm, &domain.size); - return domain; + int size = 1; + MPI_Comm_size(communicator_of(communicator), &size); + return size; #else TORCH_CHECK(false, "latent_charge_sum: an evaluation across MPI ranks " diff --git a/source/op/pt/long_range/domain_exchange.h b/source/op/pt/long_range/domain_exchange.h index f423974f26..0ebd9698d4 100644 --- a/source/op/pt/long_range/domain_exchange.h +++ b/source/op/pt/long_range/domain_exchange.h @@ -19,14 +19,8 @@ namespace deepmd { namespace long_range { -struct Domain { - int rank; - int size; -}; - -// The rank of this process and the number of ranks of ``communicator``; one -// rank for a zero address. -Domain domain_of(const torch::Tensor& communicator); +// The number of ranks of ``communicator``; one for a zero address. +int domain_count(const torch::Tensor& communicator); // Concatenates, in rank order, the rows of every rank. ``rows`` is a // contiguous float64 matrix of shape (M_r, W) on the CUDA device, with the diff --git a/source/op/pt/long_range/frame_sums.cuh b/source/op/pt/long_range/frame_sums.cuh index d3f8d40af6..967277d1bc 100644 --- a/source/op/pt/long_range/frame_sums.cuh +++ b/source/op/pt/long_range/frame_sums.cuh @@ -115,6 +115,32 @@ constexpr int kFrameThreads = 256; constexpr int kSliceNodes = 2048; constexpr int kMaxSlices = 256; +// The slices of every frame of a pass over ``nodes`` nodes of ``frame_count`` +// frames, from their average count per frame. +inline int frame_slices(int64_t nodes, int64_t frame_count) { + const int64_t average = (nodes + frame_count - 1) / frame_count; + return static_cast(std::min( + kMaxSlices, + std::max(1, (average + kSliceNodes - 1) / kSliceNodes))); +} + +// The offsets of the frames of the owners on the axis of the owners: those of +// the frames, ``frame_ptr``, when every node is an owner, else those of the +// one frame that the owners share, {0, owner_count}, which ``storage`` +// receives on the device of ``options``. +inline const int64_t* owner_frame_offsets(bool every_node, + const int64_t* frame_ptr, + int owner_count, + const torch::TensorOptions& options, + torch::Tensor& storage) { + if (every_node) { + return frame_ptr; + } + storage = torch::full({2}, int64_t{owner_count}, options); + storage.select(0, 0).zero_(); + return storage.data_ptr(); +} + struct NodeRange { int64_t begin; int64_t end; diff --git a/source/op/pt/long_range/latent_charge_sum.cu b/source/op/pt/long_range/latent_charge_sum.cu index 91af0c09b5..dea9bb1681 100644 --- a/source/op/pt/long_range/latent_charge_sum.cu +++ b/source/op/pt/long_range/latent_charge_sum.cu @@ -24,15 +24,15 @@ namespace { using deepmd::long_range::ChargeSumOutputs; - -constexpr int kMaxChannels = 8; -constexpr int kMaxGaussians = 32; - -enum KernelKind : int { kLes = 0, kSog = 1 }; +using deepmd::long_range::kLes; +using deepmd::long_range::kMaxChannels; +using deepmd::long_range::kMaxGaussians; +using deepmd::long_range::kSog; // The evaluation of the frames, by the direct or the mesh sum, of the owner -// nodes [owner_begin, owner_begin + owner_count). ``box`` is defined for -// periodic frames only. +// nodes [owner_begin, owner_begin + owner_count), which are every node or lie +// in a single frame. ``box`` is defined for periodic frames only. Without +// nodes or frames the outputs are zero. struct Evaluation { const torch::Tensor& box; const torch::Tensor& coefficient; @@ -51,6 +51,18 @@ struct Evaluation { int owner_begin, int owner_count, const torch::Tensor* communicator) const { + const int64_t frame_count = n_node_per_frame.size(0); + if (coord.size(0) == 0 || frame_count == 0) { + return {torch::zeros({owner_count, 1}, coord.options()), + torch::zeros({owner_count, charge.size(1)}, charge.options()), + torch::zeros({owner_count, 3}, charge.options()), + torch::zeros({owner_count, 9}, charge.options()), + torch::zeros({frame_count, 9}, charge.options())}; + } + TORCH_CHECK( + (owner_begin == 0 && owner_count == coord.size(0)) || frame_count == 1, + "latent_charge_sum: owners other than every node lie in one " + "frame"); if (periodic) { return deepmd::long_range::mesh_sum( coord, charge, node_mask, n_node_per_frame, total_charge, box, @@ -225,7 +237,7 @@ ChargeSumOutputs latent_charge_sum(torch::Tensor coord, coefficient = coefficient.contiguous(); const Evaluation evaluation{cells, coefficient, kernel, channel, prefactor, order, spacing, periodic}; - if (deepmd::long_range::domain_of(communicator).size == 1) { + if (deepmd::long_range::domain_count(communicator) == 1) { return evaluation(coord, charge, node_mask, n_node_per_frame, total_charge, 0, static_cast(coord.size(0)), nullptr); } diff --git a/source/op/pt/long_range/mesh_sum.cu b/source/op/pt/long_range/mesh_sum.cu index 8bccbcbf90..a81773b63e 100644 --- a/source/op/pt/long_range/mesh_sum.cu +++ b/source/op/pt/long_range/mesh_sum.cu @@ -53,7 +53,7 @@ #include #include #include -#include +#include #include #include @@ -75,10 +75,11 @@ namespace { constexpr int kThreads = 256; -constexpr int kMaxGaussians = 32; constexpr int kMaxOrder = 12; -enum KernelKind : int { kLes = 0, kSog = 1 }; +using deepmd::long_range::kLes; +using deepmd::long_range::kMaxGaussians; +using deepmd::long_range::kSog; // Resolution of the fixed-point charge mesh: 2^-40 e, far below the float32 // precision of the raw charges, for a sum of at most 2^23 e at a point. @@ -1298,17 +1299,7 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, kernel == kSog ? static_cast(coefficient.numel() / 2) : 0; const auto double_options = coord.options(); const auto float_options = charge.options(); - if (node_count == 0 || frame_count == 0) { - return {torch::zeros({owner_count, 1}, double_options), - torch::zeros({owner_count, channels}, float_options), - torch::zeros({owner_count, 3}, float_options), - torch::zeros({owner_count, 9}, float_options), - torch::zeros({frame_count, 9}, float_options)}; - } const bool every_node = owner_begin == 0 && owner_count == node_count; - TORCH_CHECK(every_node || frame_count == 1, - "latent_charge_sum: owners other than every node lie in one " - "frame"); const auto stream = at::cuda::getCurrentCUDAStream(); // The device reads pinned cells in place, and a copy of pageable ones; a // CUDA graph that captures the call reads the cells again at every replay, @@ -1341,15 +1332,8 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, } // === Step 2. The buffers of the call === - const int64_t average_nodes = (node_count + frame_count - 1) / frame_count; - const auto slices_for = [](int64_t nodes) { - return static_cast(std::min( - kMaxSlices, - std::max(1, (nodes + kSliceNodes - 1) / kSliceNodes))); - }; - const int slice_count = slices_for(average_nodes); - const int owner_slice_count = slices_for( - (int64_t{owner_count} + frame_count - 1) / std::max(frame_count, 1)); + const int slice_count = frame_slices(node_count, frame_count); + const int owner_slice_count = frame_slices(owner_count, frame_count); const int64_t stencil_bytes = dispatch_order_value(order, [](auto constant) { return static_cast(sizeof(AxisStencil)); }); @@ -1508,16 +1492,10 @@ ChargeSumOutputs mesh_sum(const torch::Tensor& coord, }); C10_CUDA_KERNEL_LAUNCH_CHECK(); } - // The frames of the owners, on the axis of the owners: the frames - // themselves when every node is an owner, else the one frame they share. - torch::Tensor owner_frame_tensor; - const int64_t* owner_frame_ptr = frame_ptr; - if (!every_node) { - owner_frame_tensor = - torch::full({2}, int64_t{owner_count}, counts.options()); - owner_frame_tensor.select(0, 0).zero_(); - owner_frame_ptr = owner_frame_tensor.data_ptr(); - } + torch::Tensor owner_frame_storage; + const int64_t* owner_frame_ptr = + owner_frame_offsets(every_node, frame_ptr, owner_count, counts.options(), + owner_frame_storage); const bool* owner_mask = node_mask.data_ptr() + owner_begin; const double* owner_coord = coord.data_ptr() + 3 * owner_begin; const float* owner_charge = node_charge + int64_t{owner_begin} * channels; diff --git a/source/tests/infer/gen_dpa4c_long_range.py b/source/tests/infer/gen_dpa4c_long_range.py index bb33fb39c8..39357e7e73 100644 --- a/source/tests/infer/gen_dpa4c_long_range.py +++ b/source/tests/infer/gen_dpa4c_long_range.py @@ -157,7 +157,9 @@ def _check_metadata(pt2_path: str, lower: str, periodic: bool = False) -> None: assert md["dim_chg_spin"] == 2 else: # The compressed descriptor folds the charge state; the archive - # carries the rebuild of every constant that holds it. + # carries the rebuild of every constant that holds it: the four + # charge-state artifacts of the descriptor and the total charge of + # the constraint. assert md["dim_chg_spin"] == 0 assert len(md["charge_state_constants"]) == 5 # The fused long-range operator serves the charge network of one diff --git a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py index 32748da63e..4f801e328c 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa1_cuda.py @@ -1188,7 +1188,7 @@ def test_fused_network_matches_the_layer_kernels_bitwise(self) -> None: (256 * sms + 37, 256 * sms), ) gen = torch.Generator(device=self.device).manual_seed(11) - for width in (96, 192, 256, 384): + for width in (64, 96, 192, 256, 384): networks = ( (70, 3, True), (2 * width + 8, 3, True), diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py index f44525f15e..a23f8713a6 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c_cuda.py @@ -3745,7 +3745,7 @@ def _fitting_kernels(call: Callable[[], Any]) -> list[str]: @_GPU -@pytest.mark.parametrize("width", [96, 192, 256, 384]) # fused hidden width +@pytest.mark.parametrize("width", [64, 96, 192, 256, 384]) # fused hidden width @pytest.mark.parametrize("run", sorted(_ROUTE_RUNS)) # run size regime def test_every_fitting_route_returns_the_bits_of_the_rules( width: int, run: str diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 0c69d53883..33680cba81 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -2,6 +2,14 @@ """Long-range latent-charge modules on the DPA4C graph lower (pt_expt).""" import copy +import json +import os +import pathlib +import subprocess +import sys +from typing import ( + Any, +) import numpy as np import pytest @@ -1605,12 +1613,13 @@ def _canonical_inputs( ) -# The layouts of the fused parity test: the width of the energy fitting, which -# selects the per-layer kernels at 32 and the fused network at 96, the hidden -# layers of the charge network, and the frames of the batch. +# The layouts of the fused parity test: the width of the energy fitting and the +# hidden layers of the charge network, which select the per-layer kernels at +# widths of 32 and 16 and the fused network at widths of 96 and 64, and the +# frames of the batch. FUSED_LAYOUTS = { - "per_layer": {"width": 32, "charge": [16, 16], "frames": 2}, - "fused": {"width": 96, "charge": [16, 16, 16], "frames": 1}, + "per_layer": {"width": 32, "charge": [16, 16, 16], "frames": 2}, + "fused": {"width": 96, "charge": [64, 64], "frames": 1}, } @@ -1634,10 +1643,11 @@ def test_fused_operator_matches_composition( With a charge network of one channel, the compact canonical lower evaluates the descriptor, both networks, the long range and the force assembly in one operator; the composition of operators, which serves the - other charge networks, is the reference. The energy fitting takes the - per-layer kernels or the fused network, over the whole node axis or in - runs of 16 nodes with a shorter last one, with and without the per-atom - virial, on two frames, the second with fewer real atoms, or on one. + other charge networks, is the reference. The energy fitting and the + charge network take the per-layer kernels or the fused network, over the + whole node axis or in runs of 16 nodes with a shorter last one, with and + without the per-atom virial, on two frames, the second with fewer real + atoms, or on one. """ from deepmd.pt_expt.kernels.dpa4c import ( canonical_long_range, @@ -1685,12 +1695,21 @@ def test_fused_operator_matches_composition( model, coord, atype, box if periodic else None ) # The rules of the fitting routes take the network of the layout on every - # run of the energy fitting; a device on which they send the runs of the - # fused layout to the per-layer kernels has no fused network to compare. - weights = fitting_operator_arguments(model.atomic_model.fitting_net).weights + # run of the energy fitting and of the charge network; a device on which + # they send the runs of the fused layout to the per-layer kernels has no + # fused network to compare. nodes = atype_t.shape[0] runs = {nodes} if tile == 0 else {tile, nodes % tile or tile} - routes = [int(torch.ops.deepmd.graph_fitting_route_rule(weights, n)) for n in runs] + networks = ( + model.atomic_model.fitting_net, + model.atomic_model.long_range.charge_net, + ) + weights = [fitting_operator_arguments(network).weights for network in networks] + routes = [ + int(torch.ops.deepmd.graph_fitting_route_rule(ws, n)) + for ws in weights + for n in runs + ] if layout == "fused" and -1 in routes: pytest.skip("the fitting rules take the per-layer kernels on this device") assert all((route == -1) == (layout == "per_layer") for route in routes), routes @@ -1728,6 +1747,354 @@ def evaluate(do_atomic_virial: bool) -> dict[str, torch.Tensor]: assert error <= 1e-5 * scale + 1e-6, (key, error.item(), scale.item()) +#: Bit widths of the template parameters of a fused fitting shape in its route +#: key (FusedShape::key of graph_fitting.cu), from the lowest bit up. +_SHAPE_KEY_FIELDS = (9, 6, 6, 6, 6, 6, 6, 6, 7, 4) + + +def _fused_shape_name(key: int) -> str: + """Return the template arguments of the fused shape of a route key, as a trace names them.""" + values = [] + for bits in _SHAPE_KEY_FIELDS: + values.append(key & ((1 << bits) - 1)) + key >>= bits + return "FusedShape<" + ", ".join(str(v) for v in values) + ">" + + +def _network_routes(regime: str) -> dict: + """Evaluate the fused long-range operator under every route of either network. + + Builds a compressed model with an energy fitting of width 96 and a charge + network of width 64 on a cluster of a run of fewer 32-node panels than + multiprocessors (``latency``) or of more (``throughput``), and evaluates + the operator under empty tables, the rules, under a table naming each + route of one network at a time, under the table that the route tuner holds + for the charge network, as the energy model looks it up, here a route + besides its rule, and, split into a full run and a shorter tail, under + tables that send one of them to the per-layer kernels of the charge + network and the other to a fused shape, and both to fused shapes. The + fitting kernels of an evaluation come from a kernel trace. + + Returns + ------- + dict + ``rules``, the fitting kernels of the rules; ``routes``, for every + untiled table, whether all six outputs equal those of the rules bit + for bit and its fitting kernels; ``mixed``, the same for every mixed + table against the rules tiled alike, with its launches of the + per-layer forward kernels; and ``fused_layers``, those launches with + both runs of the charge network on fused shapes, or ``None`` where + either run has no fused shape. + """ + from torch.profiler import ( + ProfilerActivity, + profile, + ) + + from deepmd.pt_expt.kernels import ( + fitting_routes, + ) + from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + dpa4c_canonical_compress_long_range_energy_force, + ) + from deepmd.pt_expt.kernels.fitting_routes import ( + PER_LAYER, + ROUTE_TABLE_MAGIC, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_operator_arguments, + ) + from deepmd.pt_expt.model.make_model import ( + fused_atom_bias, + ) + + config = _config("les", constraint={"source": "neutral"}) + config["descriptor"].update(n_radial=8, precision="float32") + config["fitting_net"].update( + neuron=[96] * 2, activation_function="silu", resnet_dt=False + ) + config["fitting_net"]["precision"] = "float32" + config["long_range"].update(channels=1, charge_net={"neuron": [64, 64], "seed": 23}) + model = get_model(config).to(env.DEVICE).eval() + generator = torch.Generator().manual_seed(5) + with torch.no_grad(): + for parameter in model.parameters(): + noise = torch.randn(parameter.shape, generator=generator) + parameter.add_(0.05 * noise.to(parameter)) + model.atomic_model.enable_compression(0.5) + device = env.DEVICE.index or 0 + multiprocessors = torch.cuda.get_device_properties(env.DEVICE).multi_processor_count + nodes = 30 if regime == "latency" else int(32 * multiprocessors * 1.5) + rng = np.random.default_rng(3) + # A cluster of about one atom per 20 cubic angstroms. + side = (20.0 * nodes) ** (1.0 / 3.0) + coord = rng.uniform(0.0, side, (1, nodes, 3)) + 100.0 + atype = rng.integers(0, 2, (1, nodes)) + atype_t, graph, coord_t, _ = _canonical_inputs(model, coord, atype, None) + learned, _ = model.atomic_model.fused_decomposition() + atom_bias = fused_atom_bias(model.atomic_model, learned) + fingerprint = [ + int(v) for v in torch.ops.deepmd.graph_fitting_route_fingerprint(device) + ] + + def evaluate(routes: list[int], charge_routes: list[int]) -> tuple: + return dpa4c_canonical_compress_long_range_energy_force( + learned.descriptor, + learned.fitting_net, + learned.long_range, + graph, + atype_t, + coord_t, + None, + None, + atom_bias, + fitting_routes=routes, + charge_fitting_routes=charge_routes, + ) + + def choices(width: int, size: int = nodes) -> list[int]: + fused = torch.ops.deepmd.graph_fitting_route_choices(device, width, size) + return [PER_LAYER, *(int(key) for key in fused)] + + def kernel_names(call: Any) -> list[str]: + with profile(activities=[ProfilerActivity.CUDA]) as trace: + call() + torch.cuda.synchronize() + names = [ + event.name + for event in trace.events() + if event.device_type == torch.autograd.DeviceType.CUDA + ] + assert names, "the profiler recorded no CUDA kernel" + return names + + def kernels(call: Any) -> list[str]: + names = kernel_names(call) + return sorted( + {name for name in names if "fused_network_kernel" in name} + | {"per-layer" for name in names if "layer_forward_kernel" in name} + ) + + def layer_launches(call: Any) -> int: + return sum("layer_forward_kernel" in name for name in kernel_names(call)) + + results = {} + with torch.no_grad(): + rules = evaluate([], []) + tables = [ + ("energy", c, [ROUTE_TABLE_MAGIC, *fingerprint, 1, c], []) + for c in choices(96) + ] + tables += [ + ("charge", c, [], [ROUTE_TABLE_MAGIC, *fingerprint, 1, c]) + for c in choices(64) + ] + charge_net = learned.long_range.charge_net + rule = int( + torch.ops.deepmd.graph_fitting_route_rule( + fitting_operator_arguments(charge_net).weights, nodes + ) + ) + measured = next(c for c in choices(64) if c != rule) + key = ( + fitting_routes._fingerprint(env.DEVICE), + fitting_routes.network_key(charge_net), + ) + table = (ROUTE_TABLE_MAGIC, *fingerprint, 1, measured) + fitting_routes._ROUTE_TABLES[key] = table + tables.append( + ("charge-measured", measured, [], fitting_routes.fitting_routes(charge_net)) + ) + for network, choice, routes, charge_routes in tables: + routed = evaluate(routes, charge_routes) + equal = all( + torch.equal(expected, actual) + for expected, actual in zip(rules, routed, strict=True) + ) + launched = kernels(lambda: evaluate(routes, charge_routes)) + results[f"{network} {choice}"] = {"equal": equal, "kernels": launched} + rules_kernels = kernels(lambda: evaluate([], [])) + # One full run and a shorter tail, of which one takes the per-layer + # kernels of the charge network and the other a fused shape, against + # both on fused shapes. + tile = nodes // 2 + 1 + tail = nodes - tile + os.environ["DP_NODE_TILE"] = str(tile) + tiled_rules = evaluate([], []) + fused_tail, fused_full = choices(64, tail)[1:2], choices(64, tile)[1:2] + mixed = {} + fused_layers = None + + def two_runs(tail_route: int, full_route: int) -> list[int]: + return [ROUTE_TABLE_MAGIC, *fingerprint, 2, tail, tail_route, full_route] + + if fused_tail and fused_full: + # The per-layer kernels of the energy fitting alone. + table = two_runs(*fused_tail, *fused_full) + fused_layers = layer_launches(lambda: evaluate([], table)) + pairs = [(PER_LAYER, *fused_full), (*fused_tail, PER_LAYER)] + for tail_route, full_route in pairs: + table = two_runs(tail_route, full_route) + routed = evaluate([], table) + equal = all( + torch.equal(expected, actual) + for expected, actual in zip(tiled_rules, routed, strict=True) + ) + mixed[f"{tail_route} {full_route}"] = { + "equal": equal, + "kernels": kernels(lambda: evaluate([], table)), + "layers": layer_launches(lambda: evaluate([], table)), + } + del os.environ["DP_NODE_TILE"] + return { + "rules": rules_kernels, + "routes": results, + "mixed": mixed, + "fused_layers": fused_layers, + } + + +@pytest.mark.skipif( + env.DEVICE.type != "cuda", reason="the compact canonical lower is CUDA only" +) +@pytest.mark.parametrize("regime", ["latency", "throughput"]) # run size regime +def test_every_network_route_returns_the_bits_of_the_rules(regime: str) -> None: + """Every route of either network gives the outputs of the rules bit for bit. + + The fused long-range operator evaluates the energy fitting and the charge + network of a module by a fitting-route table each. A table that names the + per-layer kernels or a fused shape for one network routes that network + there and leaves the other at its rules, which the fitting kernels of a + trace show, and every route evaluates the same contractions in the same + order, so all six outputs equal those of the rules, which empty tables + take. The table that the route tuner holds for the charge network, as the + energy model looks it up, routes it as well, and a tiled call whose full + run and tail take the per-layer kernels and a fused shape of the charge + network, in either order, keeps the bits of the rules. The trace runs in a + process of its own, untiled but for that call, since CUPTI can stop + delivering kernel records late in a long process. A run of fewer 32-node + panels than multiprocessors and one of more cover the latency and the + throughput shapes. + """ + from deepmd.pt_expt.kernels.fitting_routes import ( + PER_LAYER, + ) + + root = pathlib.Path(__file__).parents[__name__.count(".")] + code = ( + "import importlib, json, sys\n" + f"sys.path[:0] = {[str(root), *sys.path]!r}\n" + f"module = importlib.import_module({__name__!r})\n" + f"print('result:', json.dumps(module._network_routes({regime!r})))\n" + ) + process = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + check=False, + env={**os.environ, "DP_NODE_TILE": "0"}, + ) + assert process.returncode == 0, process.stdout + "\n" + process.stderr + (report,) = [ + line for line in process.stdout.splitlines() if line.startswith("result:") + ] + report = json.loads(report.removeprefix("result:")) + rules, results = report["rules"], report["routes"] + fused = [label for label in results if int(label.split()[1]) != PER_LAYER] + assert any(label.startswith("charge") for label in fused), ( + "no fused shape of the charge network runs on this device" + ) + + def shapes(kernels: list[str], width: int) -> list[str]: + return [k for k in kernels if f"FusedShape<{width}," in k] + + for label, result in results.items(): + network, choice = label.split()[0], int(label.split()[1]) + assert result["equal"], label + width, other = (96, 64) if network == "energy" else (64, 96) + own = shapes(result["kernels"], width) + if choice == PER_LAYER: + assert not own and "per-layer" in result["kernels"], (label, result) + else: + shape = _fused_shape_name(choice) + assert own and all(shape in k for k in own), (label, result) + assert shapes(result["kernels"], other) == shapes(rules, other), ( + label, + result, + rules, + ) + assert report["mixed"], "no fused shape of the charge network serves both runs" + for label, result in report["mixed"].items(): + assert result["equal"], label + assert shapes(result["kernels"], 64), (label, result) + # The run of the per-layer kernels of the charge network adds layer + # launches to those of the energy fitting. + assert result["layers"] > report["fused_layers"], (label, result, report) + + +@pytest.mark.skipif( + env.DEVICE.type != "cuda", reason="the compact canonical lower is CUDA only" +) +@pytest.mark.parametrize("charge_width", [64, 96]) # charge network width +def test_the_route_tuner_measures_the_charge_network( + monkeypatch, charge_width: int +) -> None: + """The route tuner measures the charge network of a fused module apart. + + The tuner records a table for the energy fitting and one for the charge + network, each under its own network key, which the energy model passes to + the fused long-range operator; a charge network of the shape of the + energy fitting shares its key and its table, measured once. The + measurement itself is replaced by a stand-in that records which network + it times. + """ + from deepmd.pt_expt.kernels import ( + fitting_routes, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_operator_arguments, + ) + + config = _config("les", constraint={"source": "neutral"}) + config["descriptor"].update(n_radial=8, precision="float32") + config["fitting_net"].update( + neuron=[96] * 2, activation_function="silu", resnet_dt=False + ) + config["fitting_net"]["precision"] = "float32" + config["long_range"].update( + channels=1, + charge_net={"neuron": [charge_width] * 2, "seed": 23}, + ) + config["long_range"]["charge_net"].update( + activation_function="silu", resnet_dt=False + ) + model = get_model(config).to(env.DEVICE).eval() + model.atomic_model.enable_compression(0.5) + learned, _ = model.atomic_model.fused_decomposition() + energy_key = fitting_routes.network_key(learned.fitting_net) + charge_key = fitting_routes.network_key(learned.long_range.charge_net) + assert energy_key is not None and charge_key is not None + assert (energy_key == charge_key) == (charge_width == 96) + + measured = [] + + def measure(bench, ws, largest): + measured.append(int(ws[0].shape[1])) + return [], [fitting_routes.RULES] + + tables: dict = {} + monkeypatch.setattr(fitting_routes, "_ROUTE_TABLES", tables) + monkeypatch.setattr(fitting_routes, "_measure", measure) + device = torch.device("cuda", env.DEVICE.index or 0) + fitting_routes.tune_fitting_routes(model, 0, device) + fingerprint = fitting_routes._fingerprint(device) + assert set(tables) == {(fingerprint, energy_key), (fingerprint, charge_key)} + assert sorted(measured) == sorted({96, charge_width}) + for network in (learned.fitting_net, learned.long_range.charge_net): + widths = [int(w.shape[1]) for w in fitting_operator_arguments(network).weights] + assert widths[0] in measured + + def test_library_without_the_fused_operator_keeps_the_composition(monkeypatch) -> None: """A library without the fused long-range operator evaluates the composition. From 4a5eaead03c0c80d0450624492cf455abf95d68c Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 30 Oct 2026 19:23:23 +0800 Subject: [PATCH 04/15] fix(long-range): keep the self term of a periodic frame without wave numbers The explicit k-grid evaluation subtracted the self term k_e K(0) q only in a frame that kept a wave number within the cutoff. A small cell with a wide kernel keeps none, so its lattice sum reduces to the self term alone, which the evaluation dropped: a 3 A cell with an LES kernel of alpha 0.03 gave zero energy where the mesh evaluation and the closed form give -0.2437 eV per atom, with wrong charge derivatives as well. Every node of a periodic frame now carries its self term; the k grid of a non-periodic batch holds no nodes and still gives zero. --- deepmd/dpmodel/long_range/reciprocal.py | 14 +++--- .../tests/common/dpmodel/test_long_range.py | 47 +++++++++++++++++++ 2 files changed, 54 insertions(+), 7 deletions(-) diff --git a/deepmd/dpmodel/long_range/reciprocal.py b/deepmd/dpmodel/long_range/reciprocal.py index 27b74e00e4..09a574e46e 100644 --- a/deepmd/dpmodel/long_range/reciprocal.py +++ b/deepmd/dpmodel/long_range/reciprocal.py @@ -125,8 +125,9 @@ def kspace_energy( Returns ------- Array - The energy of every node with shape ``(N,)``, in eV; zero in a frame - without wave numbers. + The energy of every node with shape ``(N,)``, in eV. A periodic frame + whose wave numbers all lie beyond the cutoff keeps its self term; a + non-periodic batch, whose k grid holds no nodes, has zero energy. """ xp = array_api_compat.array_namespace(charge) device = array_api_compat.device(charge) @@ -164,12 +165,11 @@ def kspace_energy( xp.matrix_transpose(sin_phase), weight * rho_imag ) potential = xp.sum(block_potential, axis=1) - # The self term belongs to the lattice sum of a frame with wave numbers; - # a non-periodic batch evaluated by a periodic module carries none. - has_points = xp.astype(xp.any(k_grid.mask, axis=(1, 2)), dtype)[:, None, None] + # Every node of a periodic frame carries its self term, whether or not a + # wave number of the frame lies within the cutoff; the slots of a + # non-periodic batch hold no nodes and hence no charge. potential = ( - potential - - COULOMB_CONSTANT * kernel.contact_value(charge) * charge_slot * has_points + potential - COULOMB_CONSTANT * kernel.contact_value(charge) * charge_slot ) slot_energy = 0.5 * xp.sum(charge_slot * potential, axis=-1) # === Step 5. Energy on the node axis === diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index 2ba5dffb62..be2c295ba1 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -431,6 +431,53 @@ def test_kspace_energy_of_padded_frames_matches_frames_alone( start += n +def test_kspace_energy_keeps_the_self_term_without_wave_numbers() -> None: + """A periodic frame without wave numbers below the cutoff keeps its self term. + + For a small cell and a wide kernel, every wave number of the cell lies + beyond the cutoff, so the lattice sum reduces to the self term + -k_e K(0) q_i^2 / 2 of every node, which the mesh evaluation gives as + well, over both the ragged and the fixed k grid; a non-periodic batch has + zero energy. + """ + from deepmd.dpmodel.long_range.reciprocal import ( + kspace_energy, + mesh_energy, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_mesh_grid, + ) + from deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph import ( + build_fixed_k_grid, + empty_k_grid, + ) + + kernel = LESKernel(alpha=0.03) + k_max = kernel.kspace_cutoff() + box = np.diag([3.0, 3.0, 3.0])[None] + assert 2.0 * np.pi / 3.0 > k_max + coord = np.array([[0.2, 0.3, 0.4], [1.7, 1.6, 1.5]]) + charge = np.array([[1.0], [-1.0]]) + n_node = np.array([2]) + expected = -0.5 * COULOMB_CONSTANT * 2.0 * 0.03 / np.sqrt(np.pi) * charge[:, 0] ** 2 + for grid in ( + build_k_grid(box, k_max, coord, n_node), + build_fixed_k_grid(box, np.array([1, 1, 1]), k_max, coord, n_node), + ): + assert not np.any(grid.mask) + np.testing.assert_allclose( + kspace_energy(charge, grid, kernel), expected, rtol=1e-12 + ) + mesh = build_mesh_grid(box, 0.5, coord) + frame_id = np.zeros(2, dtype=np.int64) + np.testing.assert_allclose( + mesh_energy(charge, mesh, frame_id, kernel, 8), expected, rtol=1e-4 + ) + np.testing.assert_array_equal( + kspace_energy(charge, empty_k_grid(1, coord), kernel), 0.0 + ) + + # --------------------------------------------------------------------------- # Fixed k-grid mode # --------------------------------------------------------------------------- From 5f4dd240be9fdeca4a2366052f5a136374fb8453 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 30 Oct 2026 19:23:26 +0800 Subject: [PATCH 05/15] fix(long-range): reject a composition without the NeighborGraph lower A linear_ener composition evaluates on the dense route as soon as one sub-model lacks the NeighborGraph lower, such as a pairtab sub-model, and a long-range module is evaluated on that lower only. Such a composition with a long-range sub-model was built and then failed at its first evaluation; it is now rejected when it is built, as a standalone model with such a descriptor is. --- .../atomic_model/linear_atomic_model.py | 23 +++++++++----- .../tests/common/dpmodel/test_long_range.py | 31 +++++++++++++++++++ 2 files changed, 46 insertions(+), 8 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index e0476272a3..0d692fe0a5 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -118,18 +118,25 @@ def __init__( + ", ".join(f"{type(m).__name__}={dim_of(m)}" for m in models) ) - # A long-range module attributes each pair energy to both atoms of - # the pair. It is not combined with a bridging term, which defines the - # close-pair limit of a learned local energy, nor with atom exclusion, - # which removes atomic outputs. - if any(isinstance(m, InnerPotentialAtomicModel) for m in models) and any( - m.has_long_range() for m in models - ): + # A long-range module is evaluated on the NeighborGraph lower and + # attributes each pair energy to both atoms of the pair. It is not + # combined with a sub-model without that lower, which sends the whole + # composition to the dense route, with a bridging term, which defines + # the close-pair limit of a learned local energy, nor with atom + # exclusion, which removes atomic outputs. + long_range = any(m.has_long_range() for m in models) + if long_range and not all(m.uses_graph_lower() for m in models): + raise NotImplementedError( + "a sub-model carrying a long-range module is evaluated on the " + "NeighborGraph lower, which every sub-model of the composition " + "must provide" + ) + if long_range and any(isinstance(m, InnerPotentialAtomicModel) for m in models): raise NotImplementedError( "an inner_potential bridging term cannot be combined with a " "sub-model carrying a long-range module" ) - if self.atom_exclude_types and any(m.has_long_range() for m in models): + if self.atom_exclude_types and long_range: raise NotImplementedError( "`atom_exclude_types` is not supported with a sub-model carrying " "a long-range module" diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index be2c295ba1..405b701d4e 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -2,6 +2,7 @@ """Primitives of the long-range modules: pair topology, kernels and constraint.""" import itertools +import pathlib import numpy as np import pytest @@ -176,6 +177,36 @@ def test_kernel_defaults_to_les() -> None: np.testing.assert_allclose(np.exp(kernel.log_alpha), 1.0) +def test_composition_without_the_graph_lower_is_rejected( + tmp_path: pathlib.Path, +) -> None: + """A long-range sub-model takes no sibling without the NeighborGraph lower. + + A ``pairtab`` sub-model has no NeighborGraph lower, so the composition + would evaluate on the dense route, which a long-range module does not + serve; the composition is rejected when it is built. + """ + table = tmp_path / "table.txt" + np.savetxt(table, np.column_stack([np.linspace(0.0, 6.0, 61), *[np.zeros(61)] * 3])) + child = { + "type": "standard", + "type_map": ["A", "B"], + "descriptor": {"type": "dpa4c", "rcut": 3.0, "channels": 8, "seed": 1}, + "fitting_net": {"type": "ener", "neuron": [8], "seed": 2}, + "long_range": {"type": "latent_charge", "charge_net": {"neuron": [4]}}, + } + pairtab = {"type": "pairtab", "tab_file": str(table), "rcut": 5.0, "sel": 20} + with pytest.raises(NotImplementedError, match="NeighborGraph"): + get_model( + { + "type": "linear_ener", + "type_map": ["A", "B"], + "weights": "sum", + "models": [child, pairtab], + } + ) + + def _long_range(constraint: dict | None, channels: int = 2) -> LatentChargeLongRange: charge_net = LatentChargeFitting( "latent_charge", 2, 4, channels, neuron=[4], numb_fparam=1, seed=1 From 6fcbb74f62dbd8ea06b6034a9b18dd597d7224fb Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 30 Oct 2026 19:23:29 +0800 Subject: [PATCH 06/15] test(long-range): check the pair virial independently and integrate with SciPy The pair-virial test compares the frame virial with sum_i r_i (x) f_i, which holds for central pairs in a cluster and does not share the attribution of the atomic virial. The kernel transform tests integrate with scipy.integrate.trapezoid, which serves every supported NumPy. --- source/tests/common/dpmodel/test_long_range.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index 405b701d4e..6a5dfd0c0d 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -6,6 +6,9 @@ import numpy as np import pytest +from scipy.integrate import ( + trapezoid, +) from scipy.special import ( erf, ) @@ -96,6 +99,8 @@ def test_pair_force_virial_of_a_central_potential() -> None: expected[i] += d / np.linalg.norm(d) ** 3 np.testing.assert_allclose(force, expected, rtol=1e-12) np.testing.assert_allclose(atom_virial.sum(0), virial[0], rtol=1e-12) + # The virial of central pairs in a cluster is sum_i r_i (x) f_i. + np.testing.assert_allclose(virial[0], coord.T @ force, rtol=1e-12, atol=1e-15) def test_les_kernel_closed_form() -> None: @@ -285,7 +290,7 @@ def _spherical_ft( k_r = np.sum( amplitude[None, :] * np.exp(-(r[:, None] ** 2) / (2 * width**2)), -1 ) - return float(np.trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(k * r / np.pi), r)) + return float(trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(k * r / np.pi), r)) @pytest.mark.parametrize("kernel", ["les", "sog"]) @@ -305,7 +310,7 @@ def test_kernel_kspace_matches_radial_ft(kernel: str) -> None: r = np.linspace(1e-9, 60.0, 20000) k_r = (1.0 - erf(0.8 * r)) / r # erfc(0.8 r)/r expected = 4.0 * np.pi / kk**2 - float( - np.trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(kk * r / np.pi), r) + trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(kk * r / np.pi), r) ) else: expected = _spherical_ft(kernel, kk) From ed9e304c5d3ec46bb269bebdddde6fd4aeca02f9 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 30 Oct 2026 19:23:31 +0800 Subject: [PATCH 07/15] fix(long-range): form the mesh virial from the derivative of the kernel transform The influence kernel formed the virial of a mode from the logarithmic derivative of the kernel transform times the energy of the mode, and set that derivative to zero where the transform vanishes, which drops the derivative term of a signed sum of Gaussians at such a mode. The kernel transform now returns its derivative, from which the virial of a mode is formed directly; the energies are unchanged bit for bit. --- source/op/pt/long_range/mesh_sum.cu | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/source/op/pt/long_range/mesh_sum.cu b/source/op/pt/long_range/mesh_sum.cu index a81773b63e..3fa09ef8c2 100644 --- a/source/op/pt/long_range/mesh_sum.cu +++ b/source/op/pt/long_range/mesh_sum.cu @@ -522,28 +522,26 @@ __device__ void load_kernel_terms(int kernel, } } -// The transform of the kernel at k^2 and its logarithmic derivative with -// respect to k^2. +// The transform of the kernel at k^2 and its derivative with respect to k^2. __device__ __forceinline__ void kernel_transform(int kernel, const KernelTerms& terms, int gaussians, float k_sq, float& value, - float& log_slope) { + float& slope) { if (kernel == kLes) { value = 4.f * static_cast(kPi) * expf(-k_sq * terms.inv_four_alpha_sq) / k_sq; - log_slope = -terms.inv_four_alpha_sq - 1.f / k_sq; + slope = value * (-terms.inv_four_alpha_sq - 1.f / k_sq); return; } value = 0.f; - float slope = 0.f; + slope = 0.f; for (int g = 0; g < gaussians; ++g) { const float term = terms.weight[g] * expf(-terms.half_width_sq[g] * k_sq); value += term; slope -= terms.half_width_sq[g] * term; } - log_slope = value != 0.f ? slope / value : 0.f; } // The influence function of one frame on its half spectrum: multiplies every @@ -602,8 +600,8 @@ __global__ void __launch_bounds__(kThreads) const bool weightless = row_weightless || 2 * m3 == n3 || (m1 == 0 && m2 == 0 && m3 == 0); float influence = 0.f; + float influence_slope = 0.f; float k[3] = {0.f, 0.f, 0.f}; - float log_slope = 0.f; if (!weightless) { #pragma unroll for (int a = 0; a < 3; ++a) { @@ -611,21 +609,26 @@ __global__ void __launch_bounds__(kThreads) } const float k_sq = k[0] * k[0] + k[1] * k[1] + k[2] * k[2]; float transform; - kernel_transform(kernel, terms, gaussians, k_sq, transform, log_slope); + float slope; + kernel_transform(kernel, terms, gaussians, k_sq, transform, slope); influence = row_factor * transform * factor[n1 + n2 + i3]; + influence_slope = row_factor * slope * factor[n1 + n2 + i3]; } // An interior plane of the half spectrum stands for its mirror as well. const float multiplicity = m3 == 0 ? 1.f : 2.f; const int index = row * half + i3; - float energy = 0.f; + float power = 0.f; for (int c = 0; c < channels; ++c) { float2& value = spectrum[int64_t{c} * modes + index]; - energy += value.x * value.x + value.y * value.y; + power += value.x * value.x + value.y * value.y; value.x *= influence; value.y *= influence; } - energy *= 0.5f * multiplicity * influence; - const float twice = 2.f * log_slope * energy; + const float energy = power * (0.5f * multiplicity * influence); + // Twice the derivative of the energy of the mode with respect to k^2, + // formed from the derivative of the transform, which stays defined + // where the transform vanishes. + const float twice = power * (multiplicity * influence_slope); row_sum[0] += energy; row_sum[1] += energy + twice * k[0] * k[0]; row_sum[2] += energy + twice * k[1] * k[1]; From 72cf6c3582dbd2af58fe121de91063aa5f2dd2e5 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 30 Oct 2026 19:23:32 +0800 Subject: [PATCH 08/15] fix(api): bump the C API version for the cell argument and check empty ranks DP_DPA4ComputeCanonicalGraphGPU takes the cell of a periodic system, so the C API version is 32. A rank without nodes of a decomposed long-range evaluation must report no owned node and no edge storage; another count is rejected as invalid dimensions, as for any other rank. --- source/api_c/include/c_api.h | 3 ++- source/api_cc/src/DeepPotPTExpt.cc | 5 +++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index ff8903d341..6eebcd2772 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -15,7 +15,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 31 +#define DP_C_API_VERSION 32 /** * @brief Neighbor list. @@ -502,6 +502,7 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, * DP_DeepPotComputeCanonicalGraphGPU, without ``d_virial``, ``d_coord``, * ``world``, ``box`` and ``stream``; a binary built against that name * fails to load rather than call a function of another argument list. + * API version 31, which no release carries, took no ``box``. * @since API version 31 */ extern void DP_DPA4ComputeCanonicalGraphGPU( diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 9646f9794a..c92252c1e2 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -2826,6 +2826,11 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( "decomposed over several ranks takes their communicator."); } if (long_range_ && world != nullptr && nall_nodes == 0) { + // A rank without nodes owns none and stores no edge. + if (nloc != 0 || edge_storage != 0) { + throw deepmd::deepmd_exception( + "invalid compact canonical graph dimensions."); + } translate_error([&] { const c10::DeviceGuard device_guard(torch::Device(torch::kCUDA, gpu_id)); const ptexpt::CallerStreamScope stream_scope(stream, gpu_id); From b09c79a6a67f2f87f96191512c69f9e8dce6ddf1 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 30 Oct 2026 19:34:30 +0800 Subject: [PATCH 09/15] fix(api): keep the C API at version 31 Version 31 is unreleased and already reserves the device-resident entry point for the long-range module; the cell argument completes it within that version, so the C API keeps version 31 and the release carries one argument list for it. --- source/api_c/include/c_api.h | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 6eebcd2772..ff8903d341 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -15,7 +15,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 32 +#define DP_C_API_VERSION 31 /** * @brief Neighbor list. @@ -502,7 +502,6 @@ extern void DP_DeepPotComputeEdgesGPUFloat32(DP_DeepPot* dp, * DP_DeepPotComputeCanonicalGraphGPU, without ``d_virial``, ``d_coord``, * ``world``, ``box`` and ``stream``; a binary built against that name * fails to load rather than call a function of another argument list. - * API version 31, which no release carries, took no ``box``. * @since API version 31 */ extern void DP_DPA4ComputeCanonicalGraphGPU( From 3751d3b7c115263a023f138023e892c5d9abba57 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 31 Oct 2026 01:11:24 +0800 Subject: [PATCH 10/15] fix(long-range): fail fast where PyTorch breaks the mesh evaluation Two PyTorch defects break periodic long-range models in the PME mesh mode, both reproduced with plain PyTorch and fixed in later releases. The AOTInductor runtime of PyTorch 2.10 and 2.11 allocates the complex zero buffer that the backward of a real FFT pads its gradient with as a boolean tensor, on the CPU and on CUDA, so a frozen .pt2 graph lower of a periodic model fails when it runs. PyTorch 2.12 allocates it correctly. The .pt2 export now rejects a graph lower that reads the PME mesh on an older PyTorch before AOTInductor compiles it; a .pte file and a compressed model that the long-range operator evaluates are not affected. The CPU code generation of Inductor in PyTorch 2.10 to 2.12 forms a wrong virial for the compiled training graph of a PME model, while energy and forces agree with the eager model; PyTorch 2.13 on the CPU and every release that compiles the model on CUDA form it correctly. The compiled training trace now rejects a PME model on the CPU with an older PyTorch. The tests of these paths skip on the affected releases, the fixture generator skips the periodic graph-lower archive there, which the C++ and LAMMPS suites already skip when absent, and the long-range documentation states both requirements. --- deepmd/pt_expt/train/training.py | 2 + deepmd/pt_expt/utils/serialization.py | 89 ++++++++++++++++++- doc/model/long-range.md | 6 +- .../test_deeppot_dpa4c_long_range_ptexpt.cc | 3 +- source/tests/infer/gen_dpa4c_long_range.py | 21 +++-- source/tests/pt_expt/model/test_long_range.py | 81 +++++++++++++++++ 6 files changed, 192 insertions(+), 10 deletions(-) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 94b0134aae..5bcf798806 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -891,6 +891,7 @@ def _trace_and_compile_graph( build_synthetic_graph_inputs, build_synthetic_topologies, check_graph_trace_torch_version, + check_mesh_compile_torch_version, count_synthetic_graph_edges, ) @@ -909,6 +910,7 @@ def _trace_and_compile_graph( # CUDA autograd-stream limitation) would otherwise mix a CPU model with # CUDA inputs and fail only on a GPU host. _trace_device = next(model.parameters()).device + check_mesh_compile_torch_version(model, _trace_device) e_real = count_synthetic_graph_edges( model, nframes=trace_nf, diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 16172cf7b3..7c8c23d82c 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -16,6 +16,7 @@ import torch from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, flatten_topologies, ) from deepmd.dpmodel.utils.nlist import ( @@ -31,6 +32,7 @@ traverse_model_dict, ) from deepmd.pt.utils.compile_compat import ( + _torch_release, traced_output_keys, ) from deepmd.pt_expt.model.graph_lower import ( @@ -337,6 +339,91 @@ def check_graph_trace_torch_version(model: torch.nn.Module) -> None: ) +def mesh_freeze_supported() -> bool: + """Whether AOTInductor runs the mesh evaluation of a periodic graph lower. + + The graph lower of a periodic long-range model evaluates its reciprocal + sum by FFTs over the mesh of every cell, and the derivatives of the real + transforms allocate complex buffers through the proxy executor of + AOTInductor. Its runtime allocates them with their declared dtype from + PyTorch 2.12 on; older runtimes allocate booleans, which the + complex-to-complex transform of the derivative rejects when the packaged + model runs. + """ + return _torch_release() >= (2, 12) + + +def mesh_compile_supported(device: torch.device) -> bool: + """Whether ``torch.compile`` forms the virial of the mesh evaluation on ``device``. + + The CPU code generation of Inductor forms the strain derivative of the + mesh energy in a compiled training graph correctly from PyTorch 2.13 on, + the CUDA code generation in every release that compiles the model. + """ + return device.type != "cpu" or _torch_release() >= (2, 13) + + +def check_mesh_freeze_torch_version(metadata: dict) -> None: + """Fail fast before AOTInductor packages a mesh evaluation it cannot run. + + Parameters + ---------- + metadata : dict + The metadata of the traced lower, with its ``lower_input_kind`` and + its ``topologies``. + + Raises + ------ + RuntimeError + If the graph lower reads the PME mesh and :func:`mesh_freeze_supported` + does not hold. + """ + if ( + metadata.get("lower_input_kind") == "graph" + and PME_MESH in metadata.get("topologies", ()) + and not mesh_freeze_supported() + ): + raise RuntimeError( + "freezing a periodic long-range model into a .pt2 graph lower " + f"requires torch >= 2.12; found torch {torch.__version__}, whose " + "AOTInductor runtime allocates the complex buffers of the mesh " + "derivatives as booleans. Upgrade torch, freeze to a .pte file, " + "or compress the model for the long-range operator on a CUDA " + "device." + ) + + +def check_mesh_compile_torch_version( + model: torch.nn.Module, device: torch.device +) -> None: + """Fail fast before compiling a model whose virial Inductor forms wrongly. + + Parameters + ---------- + model : torch.nn.Module + The model about to be traced for compiled training. + device : torch.device + The device of the trace. + + Raises + ------ + RuntimeError + If the model reads the PME mesh and :func:`mesh_compile_supported` + does not hold on ``device``. + """ + if ( + PME_MESH in model.atomic_model.required_topologies() + and not mesh_compile_supported(device) + ): + raise RuntimeError( + "compiled training of a periodic long-range model on the PME mesh " + f"requires torch >= 2.13 on the CPU; found torch {torch.__version__}, " + "whose Inductor CPU code generation forms a wrong virial for it. " + "Upgrade torch, train on a CUDA device, or set " + "`training.enable_compile` to false." + ) + + # Module-level cache for the trace-time sendlist buffer. The pointer # value embedded in ``send_list_tensor`` references this numpy array's # data; the array must outlive the trace + export call. Caching here @@ -765,7 +852,6 @@ def build_synthetic_topologies( """ from deepmd.dpmodel.utils.neighbor_graph import ( K_GRID, - PME_MESH, KGrid, PairGraph, ) @@ -3117,6 +3203,7 @@ def _deserialize_to_file_pt2( metadata["has_comm_artifact"] = with_comm and bool( metadata.get("has_comm_artifact") ) + check_mesh_freeze_torch_version(metadata) # A charge-state update swaps constants in one compiled artifact. A model # that also needs the with-comm lower would own two independent constant diff --git a/doc/model/long-range.md b/doc/model/long-range.md index 81e673febf..8797277cea 100644 --- a/doc/model/long-range.md +++ b/doc/model/long-range.md @@ -200,11 +200,15 @@ Training uses the usual energy loss; see [training energy models](train-energy.md). With `enable_compile`, the long-range interaction is part of the compiled model: the pairs, the wave numbers or the mesh of each batch are built outside the compiled region and -enter it as inputs of dynamic size. +enter it as inputs of dynamic size. Compiled training of a periodic model in +the `pme` mode on the CPU requires PyTorch 2.13 or later. `dp --pt-expt freeze` exports the model, long-range module included, to `.pt2` or `.pte`, and `dp --pt-expt compress` compresses its descriptor into a `.pt2` archive while the networks of the long-range module stay as they are. +Freezing a model with `periodic: true` to `.pt2` requires PyTorch 2.12 or +later, except for a model that the dedicated operator described below +evaluates, whose mesh the operator transforms. A trained checkpoint (`.pt`) and the frozen model are evaluated by `dp test` and the Python inference interface; the C++ interface and LAMMPS, with `pair_style deepmd` and `pair_style deepmd/kk`, read the `.pt2` archive. diff --git a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc index 9658effbc1..337886ce94 100644 --- a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc @@ -304,7 +304,8 @@ class TestInferDeepPotDpa4cLongRangePeriodicPtExpt : public ::testing::Test { if (!model_file.good()) { GTEST_SKIP() << "Skip because " << model_path << " was not generated (gen_dpa4c_long_range.py writes the " - "compressed archives on a CUDA device only)."; + "compressed archives on a CUDA device only, and the " + "graph lower with torch >= 2.12)."; } dp.init(model_path); deepmd_test::ExpectedRef ref; diff --git a/source/tests/infer/gen_dpa4c_long_range.py b/source/tests/infer/gen_dpa4c_long_range.py index 39357e7e73..96de474c97 100644 --- a/source/tests/infer/gen_dpa4c_long_range.py +++ b/source/tests/infer/gen_dpa4c_long_range.py @@ -13,7 +13,8 @@ - ``deeppot_dpa4c_long_range.pt2`` and ``deeppot_dpa4c_long_range_periodic.pt2``: uncompressed, fp64, graph lower, - which compiles on the CPU; + which compiles on the CPU; the periodic one needs torch >= 2.12, whose + AOTInductor packages the mesh evaluation; - ``deeppot_dpa4c_long_range_compressed.pt2`` and ``deeppot_dpa4c_long_range_periodic_compressed.pt2``: compressed, fp32, compact canonical lower, whose operators are CUDA only; they are generated @@ -230,6 +231,7 @@ def _export_compressed(pt2_path: str, config: dict) -> None: def main() -> None: from deepmd.pt_expt.utils.serialization import ( deserialize_to_file, + mesh_freeze_supported, ) ensure_inductor_compiler() @@ -246,12 +248,17 @@ def main() -> None: "software": "deepmd-kit", "version": "3.0.0", } - print(f"Exporting to {stem}.pt2 (lower_kind='graph') ...") # noqa: T201 - deserialize_to_file( - f"{stem}.pt2", data, do_atomic_virial=True, lower_kind="graph" - ) - _check_metadata(f"{stem}.pt2", "graph", periodic) - _write_ref(f"{stem}.pt2", f"{stem}.expected", periodic) + if periodic and not mesh_freeze_supported(): + # The C++ and LAMMPS suites skip the archive this torch cannot + # package. + print(f"Skipping {stem}.pt2: the mesh evaluation needs torch >= 2.12") # noqa: T201 + else: + print(f"Exporting to {stem}.pt2 (lower_kind='graph') ...") # noqa: T201 + deserialize_to_file( + f"{stem}.pt2", data, do_atomic_virial=True, lower_kind="graph" + ) + _check_metadata(f"{stem}.pt2", "graph", periodic) + _write_ref(f"{stem}.pt2", f"{stem}.expected", periodic) if os.environ.get("DP_VARIANT") == "cuda": compressed = f"{stem}_compressed" diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 33680cba81..619e1d1361 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -31,6 +31,7 @@ from deepmd.pt_expt.train.training import ( _CompiledModel, _get_model_structure_key, + _trace_and_compile_graph, model_change_out_bias, ) from deepmd.pt_expt.train.wrapper import ( @@ -38,9 +39,14 @@ ) from deepmd.pt_expt.utils import ( env, + serialization, ) from deepmd.pt_expt.utils.serialization import ( + check_mesh_compile_torch_version, + check_mesh_freeze_torch_version, deserialize_to_file, + mesh_compile_supported, + mesh_freeze_supported, ) from ..compile_utils import ( @@ -1435,6 +1441,8 @@ def test_periodic_frozen_model_matches_mesh_evaluation( of a batch as inputs of dynamic size, and matches the model evaluated on the mesh; a batch without a cell keeps the direct sum. """ + if suffix == ".pt2" and not mesh_freeze_supported(): + pytest.skip(f"torch {torch.__version__} cannot package the mesh evaluation") config = _config(kernel) config["long_range"]["periodic"] = True config["long_range"]["k_grid_mode"] = grid_mode @@ -1503,6 +1511,8 @@ def test_periodic_compressed_frozen_model_matches_uncompressed( the graph lower reads the mesh. Both match the uncompressed model on the mesh, for a periodic frame and for a cluster. """ + if env.DEVICE.type != "cuda" and not mesh_freeze_supported(): + pytest.skip(f"torch {torch.__version__} cannot package the mesh evaluation") config = _config(kernel, constraint=None if source is None else {"source": source}) config["descriptor"].update(n_radial=8, precision="float32") config["fitting_net"].update( @@ -2562,6 +2572,8 @@ def test_periodic_pme_serialization_round_trip(kernel: str) -> None: @REQUIRES_SUPPORTED_COMPILE def test_periodic_pme_compiled_matches_eager(kernel: str) -> None: """The compiled graph lower evaluates the PME model like the eager one.""" + if not mesh_compile_supported(env.DEVICE): + pytest.skip(f"torch {torch.__version__} miscompiles the mesh virial here") model = _periodic_pme_model(kernel) coord, atype, box = _periodic_case() coord_t, atype_t, box_t = ( @@ -2584,3 +2596,72 @@ def test_periodic_pme_compiled_matches_eager(kernel: str) -> None: grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] assert grads and all(grad is not None for grad in grads) assert all(torch.isfinite(grad).all() for grad in grads) + + +@pytest.mark.parametrize( + ("version", "lower", "topologies", "accepted"), + [ + ("2.11.0+cpu", "graph", ["all_pairs", "pme_mesh"], False), + ("2.12.0", "graph", ["all_pairs", "pme_mesh"], True), + ("2.11.0+cpu", "dpa4c_canonical", ["all_pairs", "pme_mesh"], True), + ("2.11.0+cpu", "graph", ["all_pairs"], True), + ], +) +def test_mesh_freeze_version_guard( + monkeypatch, version: str, lower: str, topologies: list[str], accepted: bool +) -> None: + """A .pt2 graph lower that reads the PME mesh freezes from torch 2.12 on.""" + monkeypatch.setattr(torch, "__version__", version) + metadata = {"lower_input_kind": lower, "topologies": topologies} + if accepted: + check_mesh_freeze_torch_version(metadata) + else: + with pytest.raises(RuntimeError, match=r"torch >= 2\.12"): + check_mesh_freeze_torch_version(metadata) + + +def test_periodic_freeze_checks_the_torch_version(tmp_path, monkeypatch) -> None: + """The .pt2 freeze of a periodic model stops before AOTInductor where it cannot run.""" + config = _config("les") + config["long_range"]["periodic"] = True + model = get_model(config) + monkeypatch.setattr(serialization, "mesh_freeze_supported", lambda: False) + with pytest.raises(RuntimeError, match=r"torch >= 2\.12"): + deserialize_to_file( + str(tmp_path / "periodic.pt2"), + {"model": model.serialize(), "model_def_script": config}, + lower_kind="graph", + ) + + +@pytest.mark.parametrize( + ("version", "device", "grid_mode", "accepted"), + [ + ("2.12.0", "cpu", "pme", False), + ("2.13.0", "cpu", "pme", True), + ("2.11.0", "cuda", "pme", True), + ("2.12.0", "cpu", "ragged", True), + ], +) +def test_mesh_compile_version_guard( + monkeypatch, version: str, device: str, grid_mode: str, accepted: bool +) -> None: + """Compiled PME training runs on the CPU from torch 2.13 on, on CUDA throughout.""" + config = _config("les") + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = grid_mode + model = get_model(config) + monkeypatch.setattr(torch, "__version__", version) + if accepted: + check_mesh_compile_torch_version(model, torch.device(device)) + else: + with pytest.raises(RuntimeError, match=r"torch >= 2\.13"): + check_mesh_compile_torch_version(model, torch.device(device)) + + +def test_periodic_pme_compile_checks_the_torch_version(monkeypatch) -> None: + """The compiled training trace of a PME model stops where Inductor forms a wrong virial.""" + model = _periodic_pme_model("les") + monkeypatch.setattr(serialization, "mesh_compile_supported", lambda device: False) + with pytest.raises(RuntimeError, match=r"torch >= 2\.13"): + _trace_and_compile_graph(model, None, None, None, None) From ee236580a08a8dd917c203098f9c338d35aa18d4 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 31 Oct 2026 01:11:28 +0800 Subject: [PATCH 11/15] test(long-range): check the orientation of the pair virial A per-pair gradient that is not parallel to its pair makes the pair virial asymmetric, so the check of -g_p (x) pair_vec distinguishes it from its transpose, which a central potential cannot. --- source/tests/common/dpmodel/test_long_range.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index 6a5dfd0c0d..914052acde 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -103,6 +103,23 @@ def test_pair_force_virial_of_a_central_potential() -> None: np.testing.assert_allclose(virial[0], coord.T @ force, rtol=1e-12, atol=1e-15) +def test_pair_force_virial_orients_the_pair_virial() -> None: + """The virial of a pair is -g_p (x) pair_vec, the gradient index first.""" + coord = np.array([[0.0, 0.0, 0.0], [1.0, 0.2, 0.0], [0.0, 1.5, 0.3]]) + pairs = build_pair_graph(coord, np.array([3]), np.array([0, 1, 0])) + # A gradient that is not parallel to its pair, for which the virial is + # not symmetric and its orientation shows. + rng = np.random.default_rng(1) + grad = np.where(pairs.mask[:, None], rng.normal(size=pairs.vec.shape), 0.0) + _, atom_virial, virial = pair_force_virial( + grad, pairs.vec, pairs.index, pairs.mask, np.array([3]) + ) + expected = -np.einsum("pa,pb->ab", grad, pairs.vec) + assert not np.allclose(expected, expected.T) + np.testing.assert_allclose(virial[0], expected, rtol=1e-12, atol=1e-15) + np.testing.assert_allclose(atom_virial.sum(0), expected, rtol=1e-12, atol=1e-15) + + def test_les_kernel_closed_form() -> None: """LES is erf(alpha r)/r, finite at contact.""" kernel = LESKernel(alpha=0.7) From 50cb2901697d0d4d29ed06ec2632b7293d238f69 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 31 Oct 2026 16:58:57 +0800 Subject: [PATCH 12/15] fix(pt_expt): keep the lower of a frozen model through change-bias change-bias re-exported a .pte or .pt2 model through the default dense lower, which a model with a long-range module, like every graph-only model, does not implement, so the command failed for it. The updated archive now keeps the lower recorded in the input archive, as a backend conversion does, so that only the bias changes; the per-atom virial follows from the lower. --- deepmd/pt_expt/entrypoints/main.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 7bb3d6afc2..873684648b 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -851,8 +851,12 @@ def change_bias( ) ) model_dict = model_to_change.serialize() + # The updated archive keeps the lower of the input archive, as a + # backend conversion does, so that only the bias changes. deserialize_to_file( - output_path, {"model": model_dict, "model_def_script": model_params} + output_path, + {"model": model_dict, "model_def_script": model_params}, + lower_kind=pte_data["lower_input_kind"], ) log.info(f"Saved model to {output_path}") From 7c5729413fe96abbe927d6442a58bfa50bab4e37 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 31 Oct 2026 16:59:01 +0800 Subject: [PATCH 13/15] test(long-range): cover change-bias, atomic parameters and the C++ mesh topology change-bias on a frozen long-range model keeps its graph lower, shifts the energies by the change of the bias and leaves the forces unchanged. The charge network reads the atomic parameters of the energy fitting; a test separates the long-range energy from the short-range model of the same seeds and checks that it follows the atomic parameters. The C++ interface builds the mesh topology of a periodic graph lower itself. The fixture generator writes the mesh of the Python builder for the reference cell and for a frame without a cell, which needs no archive and so is written by every PyTorch release, and a C++ test compares every field that the C++ builder forms with it. --- .../test_deeppot_dpa4c_long_range_ptexpt.cc | 46 ++++++++++++++ source/tests/infer/gen_dpa4c_long_range.py | 47 +++++++++++++++ source/tests/pt_expt/model/test_long_range.py | 60 +++++++++++++++++++ 3 files changed, 153 insertions(+) diff --git a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc index 337886ce94..37b423a59a 100644 --- a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc @@ -9,6 +9,7 @@ // positions of the atoms and the cell. #include +#include #include #include #include @@ -26,6 +27,7 @@ #include #include +#include "../src/longRangePTExpt.h" #include "commonPT.h" #ifdef DEEPMD_TORCH_HAS_CUDA #include @@ -52,6 +54,8 @@ constexpr const char* kPeriodicCompressedModelPath = "../../tests/infer/deeppot_dpa4c_long_range_periodic_compressed.pt2"; constexpr const char* kPeriodicCompressedRefPath = "../../tests/infer/deeppot_dpa4c_long_range_periodic_compressed.expected"; +constexpr const char* kPeriodicMeshRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic_mesh.expected"; // The full neighbor list of a non-periodic system within ``rc``: every atom // is local, so the list holds no ghost. @@ -386,6 +390,48 @@ TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, deepmd::deepmd_exception); } +// The mesh topology that the C++ interface builds for the graph lower of a +// periodic model matches the mesh of the Python builder field by field, in +// the cell and without a cell. The reference needs no archive, so the test +// runs whatever PyTorch release generated the fixtures. +TEST(TestInferDeepPotDpa4cLongRangeMeshPtExpt, matches_python_builder) { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#else + deepmd_test::ExpectedRef ref; + ref.load(kPeriodicMeshRefPath); + // The fields in the order of the topology inputs of the graph lower. + const std::vector names = {"box", "volume", "coord", "strain", + "mask", "mode_1", "mode_2", "mode_3"}; + std::vector coord = ref.get("periodic", "coord"); + const auto coord_tensor = + torch::from_blob(coord.data(), + {static_cast(coord.size() / 3), 3}, + torch::kFloat64) + .clone(); + const double spacing = ref.get("periodic", "spacing")[0]; + for (const std::string section : {"periodic", "cluster"}) { + const std::vector box = + section == "periodic" ? ref.get(section, "box") : std::vector{}; + const auto fields = deepmd::ptexpt::buildMeshTopology( + coord_tensor, box, spacing, torch::kFloat64); + ASSERT_EQ(fields.size(), names.size()); + for (std::size_t k = 0; k < names.size(); ++k) { + const auto actual = fields[k].to(torch::kFloat64).contiguous().view(-1); + const std::vector expected = ref.get(section, names[k]); + ASSERT_EQ(actual.numel(), static_cast(expected.size())) + << section << " " << names[k]; + const double* value = actual.data_ptr(); + for (std::size_t i = 0; i < expected.size(); ++i) { + EXPECT_NEAR(value[i], expected[i], + 1e-12 * std::max(1.0, std::abs(expected[i]))) + << section << " " << names[k] << "[" << i << "]"; + } + } + } +#endif +} + #if defined(BUILD_PYTORCH) && BUILD_PT_EXPT && defined(DEEPMD_TORCH_HAS_CUDA) // The device-resident entry point of the compressed archives, which reads the // positions of the nodes and the cell of a periodic system besides the compact diff --git a/source/tests/infer/gen_dpa4c_long_range.py b/source/tests/infer/gen_dpa4c_long_range.py index 96de474c97..3c71034dd5 100644 --- a/source/tests/infer/gen_dpa4c_long_range.py +++ b/source/tests/infer/gen_dpa4c_long_range.py @@ -31,9 +31,14 @@ 9 Å apart, beyond the cutoff of the descriptor, which only the long-range module couples; in the cell, the first group meets its images across the boundary. + +``deeppot_dpa4c_long_range_periodic_mesh.expected`` holds the mesh topology +that the Python builder forms for the periodic case, against which the C++ +suite checks the mesh it builds for the graph lower; it needs no archive. """ import copy +import dataclasses import json import os import sys @@ -206,6 +211,46 @@ def _write_ref(pt2_path: str, ref_path: str, periodic: bool = False) -> None: print(f"Wrote {ref_path}") # noqa: T201 +def _write_mesh_ref(config: dict, ref_path: str) -> None: + """Write the mesh topology of the reference frame, built in Python, for the C++ twin. + + The graph lower of a periodic archive reads the mesh that the C++ + interface builds from the cell at the spacing of the archive. The + reference is the mesh of the Python builder at that spacing, in the cell + :data:`BOX` and without a cell (sections ``periodic`` and ``cluster``). + It needs no archive, so it is written by every PyTorch release. + """ + from deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph import ( + build_mesh_grid, + empty_mesh_grid, + ) + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + model = BaseModel.deserialize(_model_dict(config)) + module = model.atomic_model.long_range + module.use_mesh_evaluation() + spacing = module.topology_params()["pme_mesh"]["spacing"] + coord = _COORDS[0] + sections = {} + for name, mesh in ( + ("periodic", build_mesh_grid(BOX, spacing, coord)), + ("cluster", empty_mesh_grid(1, coord)), + ): + sections[name] = {"spacing": np.array([spacing])} + for field in dataclasses.fields(mesh): + sections[name][field.name] = np.asarray( + getattr(mesh, field.name), dtype=np.float64 + ) + write_expected_ref( + ref_path, + sections=sections, + source_script="source/tests/infer/gen_dpa4c_long_range.py", + ) + print(f"Wrote {ref_path}") # noqa: T201 + + def _export_compressed(pt2_path: str, config: dict) -> None: """Compress the fp32 model and export it for the CUDA device.""" from deepmd.pt_expt.model.model import ( @@ -248,6 +293,8 @@ def main() -> None: "software": "deepmd-kit", "version": "3.0.0", } + if periodic: + _write_mesh_ref(config, f"{stem}_mesh.expected") if periodic and not mesh_freeze_supported(): # The C++ and LAMMPS suites skip the archive this torch cannot # package. diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index 619e1d1361..a653ece4e8 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -25,6 +25,9 @@ from deepmd.infer import ( DeepPot, ) +from deepmd.pt_expt.entrypoints.main import ( + change_bias, +) from deepmd.pt_expt.model.get_model import ( get_model, ) @@ -161,6 +164,35 @@ def test_energy_matches_direct_pair_sum(kernel: str) -> None: ) +def test_charges_read_the_atomic_parameters() -> None: + """The charge network reads the atomic parameters that the energy fitting reads.""" + config = _config("les", constraint=None) + config["fitting_net"]["numb_aparam"] = 1 + short_range_config = copy.deepcopy(config) + short_range_config.pop("long_range") + model = get_model(config).to(env.DEVICE).eval() + # The same seeds give the short-range networks of both models. + short_range = get_model(short_range_config).to(env.DEVICE).eval() + neutral = copy.deepcopy(model) + with torch.no_grad(): + for parameter in neutral.atomic_model.long_range.charge_net.parameters(): + parameter.zero_() + coord, atype = _batch(seed=7) + long_range_energy = [] + for value in (0.0, 1.5): + aparam = torch.full( + (*atype.shape, 1), value, dtype=torch.float64, device=env.DEVICE + ) + energy = short_range(coord, atype, aparam=aparam)["energy"] + # Without charges the long-range model is the short-range one. + torch.testing.assert_close( + neutral(coord, atype, aparam=aparam)["energy"], energy + ) + long_range_energy.append(model(coord, atype, aparam=aparam)["energy"] - energy) + difference = (long_range_energy[1] - long_range_energy[0]).abs() + assert torch.all(difference > 1e-6), difference + + @pytest.mark.parametrize("kernel", ["les", "sog"]) def test_force_virial_hessian_match_finite_differences(kernel: str) -> None: """Force, affine-strain virial and Hessian are derivatives of the complete energy.""" @@ -870,6 +902,34 @@ def test_frozen_model_matches_direct_forward( ) +def test_change_bias_keeps_the_graph_lower(tmp_path) -> None: + """change-bias rewrites a frozen long-range model on its graph lower.""" + config = _config("les") + model = get_model(config).to(env.DEVICE).eval() + frozen = tmp_path / "long_range.pte" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + lower_kind="graph", + ) + bias = [0.5, -1.25] + updated = tmp_path / "long_range_updated.pte" + change_bias(str(frozen), bias_value=bias, output=str(updated)) + before = DeepPot(str(frozen), auto_batch_size=False) + after = DeepPot(str(updated), auto_batch_size=False) + assert after.deep_eval.metadata["lower_input_kind"] == "graph" + coord, atype = _batch(seed=7) + coord, atype = coord.cpu().numpy(), atype.cpu().numpy() + energy, force, _ = before.eval(coord, None, atype, mixed_type=True) + new_energy, new_force, _ = after.eval(coord, None, atype, mixed_type=True) + old_bias = model.get_out_bias().detach().cpu().numpy().reshape(-1) + # The padding atoms, of type -1, carry no bias. + delta = np.asarray(bias) - old_bias + shift = np.where(atype >= 0, delta[atype], 0.0).sum(axis=-1) + np.testing.assert_allclose(new_energy.reshape(-1), energy.reshape(-1) + shift) + np.testing.assert_allclose(new_force, force, rtol=1e-12, atol=1e-12) + + @pytest.mark.parametrize( ("kernel", "source"), [("les", None), ("les", "neutral"), ("sog", "charge_spin")] ) From e890b204f4a8ae7a74ad3c89b849529ec3a61f9d Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 31 Oct 2026 20:59:49 +0800 Subject: [PATCH 14/15] fix(pt_expt): carry the per-atom virial through change-bias and re-select canonical lowers change-bias kept the lower of a frozen archive but exported it without its per-atom virial, which a dense-lower archive loses that way; serialize_from_file now reports the do_atomic_virial of the archive metadata next to its lower, and change-bias passes both on. A compact canonical lower serves CUDA devices alone, so change-bias on a host without one failed in the canonical export. Such a lower is now selected anew for the host, as compression selects it: it stays canonical on a CUDA device and becomes the graph lower elsewhere. --- deepmd/pt_expt/entrypoints/main.py | 13 ++++++++++--- deepmd/pt_expt/utils/serialization.py | 5 ++++- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 873684648b..72de9277da 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -851,12 +851,19 @@ def change_bias( ) ) model_dict = model_to_change.serialize() - # The updated archive keeps the lower of the input archive, as a - # backend conversion does, so that only the bias changes. + # The updated archive keeps the lower and the per-atom virial of the + # input archive, so that only the bias changes. A compact canonical + # lower serves CUDA devices alone and is selected anew for the host, + # as compression selects it: it stays on a CUDA device, and a host + # without one exports the graph lower. + lower_kind = pte_data["lower_input_kind"] + if lower_kind in ("dpa1_canonical", "dpa4c_canonical"): + lower_kind = "auto" deserialize_to_file( output_path, {"model": model_dict, "model_def_script": model_params}, - lower_kind=pte_data["lower_input_kind"], + do_atomic_virial=pte_data["do_atomic_virial"], + lower_kind=lower_kind, ) log.info(f"Saved model to {output_path}") diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 7c8c23d82c..e5fb598494 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1618,7 +1618,8 @@ def serialize_from_file(model_file: str) -> dict: The serialized model data. If the archive contains ``model_def_script.json`` (training config), it is included under the ``"model_def_script"`` key. ``lower_input_kind`` records - the concrete lower ABI from the artifact metadata. + the concrete lower ABI and ``do_atomic_virial`` whether the archive + exports the per-atom virial, both from the artifact metadata. """ if model_file.endswith(".pt2"): return _serialize_from_file_pt2(model_file) @@ -1645,6 +1646,7 @@ def _serialize_from_file_pte(model_file: str) -> dict: model_dict["lower_input_kind"] = metadata.get( "lower_input_kind", model_dict.get("lower_input_kind", "nlist") ) + model_dict["do_atomic_virial"] = bool(metadata.get("do_atomic_virial", False)) if extra_files["model_def_script.json"]: model_dict["model_def_script"] = json.loads( extra_files["model_def_script.json"] @@ -1682,6 +1684,7 @@ def _serialize_from_file_pt2(model_file: str) -> dict: model_dict["lower_input_kind"] = metadata.get( "lower_input_kind", model_dict.get("lower_input_kind", "nlist") ) + model_dict["do_atomic_virial"] = bool(metadata.get("do_atomic_virial", False)) if model_def_script_json: model_dict["model_def_script"] = json.loads(model_def_script_json) return model_dict From 8407bdcc296d162d3c4ae71588fa0cdd5a2465b7 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 31 Oct 2026 20:59:53 +0800 Subject: [PATCH 15/15] test(long-range): cover the deployed options and the change-bias forms The frozen-model test gains a case with atomic parameters and a pair exclusion, evaluated through the deployed pair topology, and checks the lower and the per-atom virial that every archive reports. change-bias is checked on a .pt2 archive that records the compact canonical lower and on a dense-lower archive with the per-atom virial. The k-space sum is checked at a configured kspace_tolerance, the charge network against the case embedding of a multi-task model, and the SOG kernel with an amplitude shared by its Gaussians. --- .../tests/common/dpmodel/test_long_range.py | 4 + source/tests/pt_expt/model/test_long_range.py | 131 ++++++++++++++---- source/tests/pt_expt/test_change_bias.py | 23 +++ 3 files changed, 128 insertions(+), 30 deletions(-) diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py index 914052acde..543ce97dc8 100644 --- a/source/tests/common/dpmodel/test_long_range.py +++ b/source/tests/common/dpmodel/test_long_range.py @@ -137,6 +137,10 @@ def test_sog_kernel_closed_form_and_defaults() -> None: -(r**2) / (2 * 1.7**2) ) np.testing.assert_allclose(kernel(r**2), expected, rtol=1e-12) + # A scalar amplitude is shared by every Gaussian. + shared = SOGKernel(amplitude=0.4, width=[0.8, 1.7]) + expected = 0.4 * (np.exp(-(r**2) / (2 * 0.8**2)) + np.exp(-(r**2) / (2 * 1.7**2))) + np.testing.assert_allclose(shared(r**2), expected, rtol=1e-12) default = SOGKernel(n_gaussian=3, base_width=2.0, width_ratio=1.5) np.testing.assert_allclose(np.exp(default.log_width), [2.0, 3.0, 4.5]) np.testing.assert_allclose( diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py index a653ece4e8..0581dd8626 100644 --- a/source/tests/pt_expt/model/test_long_range.py +++ b/source/tests/pt_expt/model/test_long_range.py @@ -50,6 +50,7 @@ deserialize_to_file, mesh_compile_supported, mesh_freeze_supported, + serialize_from_file, ) from ..compile_utils import ( @@ -164,10 +165,18 @@ def test_energy_matches_direct_pair_sum(kernel: str) -> None: ) -def test_charges_read_the_atomic_parameters() -> None: - """The charge network reads the atomic parameters that the energy fitting reads.""" +@pytest.mark.parametrize("condition", ["aparam", "case_embd"]) +def test_charges_follow_the_fitting_conditions(condition: str) -> None: + """The charge network reads the atomic parameters and the case of the energy fitting. + + The long-range energy is the model less the short-range model of the same + seeds; it changes with the atomic parameters of the atoms and with the + case that a multi-task model selects. + """ config = _config("les", constraint=None) - config["fitting_net"]["numb_aparam"] = 1 + # Two atomic parameters, or the one-hot embedding of two cases. + key = "numb_aparam" if condition == "aparam" else "dim_case_embd" + config["fitting_net"][key] = 2 short_range_config = copy.deepcopy(config) short_range_config.pop("long_range") model = get_model(config).to(env.DEVICE).eval() @@ -179,16 +188,19 @@ def test_charges_read_the_atomic_parameters() -> None: parameter.zero_() coord, atype = _batch(seed=7) long_range_energy = [] - for value in (0.0, 1.5): - aparam = torch.full( - (*atype.shape, 1), value, dtype=torch.float64, device=env.DEVICE - ) - energy = short_range(coord, atype, aparam=aparam)["energy"] + for value in (0, 1): + kwargs = {} + if condition == "aparam": + kwargs["aparam"] = torch.full( + (*atype.shape, 2), 1.5 * value, dtype=torch.float64, device=env.DEVICE + ) + else: + for each in (model, short_range, neutral): + each.set_case_embd(value) + energy = short_range(coord, atype, **kwargs)["energy"] # Without charges the long-range model is the short-range one. - torch.testing.assert_close( - neutral(coord, atype, aparam=aparam)["energy"], energy - ) - long_range_energy.append(model(coord, atype, aparam=aparam)["energy"] - energy) + torch.testing.assert_close(neutral(coord, atype, **kwargs)["energy"], energy) + long_range_energy.append(model(coord, atype, **kwargs)["energy"] - energy) difference = (long_range_energy[1] - long_range_energy[0]).abs() assert torch.all(difference > 1e-6), difference @@ -848,18 +860,31 @@ def test_periodic_checkpoint_inference_matches_direct_forward( @pytest.mark.parametrize( - ("kernel", "source", "suffix"), - [("les", "neutral", ".pt2"), ("sog", "neutral", ".pte"), ("les", "fparam", ".pte")], + ("kernel", "source", "suffix", "atomic"), + [ + ("les", "neutral", ".pt2", False), + ("sog", "neutral", ".pte", False), + ("les", "fparam", ".pte", False), + ("les", "neutral", ".pte", True), + ], ) def test_frozen_model_matches_direct_forward( - tmp_path, kernel: str, source: str, suffix: str + tmp_path, kernel: str, source: str, suffix: str, atomic: bool ) -> None: - """A frozen model reads the pair topology as named inputs of any length.""" + """A frozen model reads the pair topology as named inputs of any length. + + With ``atomic``, the model reads atomic parameters, which reach the charge + network, and excludes the pairs of two types, which the deployed pair + topology omits. + """ config = _config(kernel, constraint={"source": source, "fparam_index": 0}) frame = {} if source == "fparam": config["fitting_net"]["numb_fparam"] = 1 frame["fparam"] = np.array([[1.0], [-0.5], [2.0]]) + if atomic: + config["fitting_net"]["numb_aparam"] = 1 + config["pair_exclude_types"] = [[0, 1]] model = get_model(config).to(env.DEVICE).eval() frozen = tmp_path / f"long_range_{kernel}{suffix}" deserialize_to_file( @@ -870,6 +895,11 @@ def test_frozen_model_matches_direct_forward( ) dp = DeepPot(str(frozen), auto_batch_size=False) assert dp.deep_eval.metadata["topologies"] == ["all_pairs"] + # The archive reports its lower and its per-atom virial, which + # change-bias carries over. + recorded = serialize_from_file(str(frozen)) + assert recorded["lower_input_kind"] == "graph" + assert recorded["do_atomic_virial"] if suffix == ".pte": # The pair topology follows the thirteen inputs of the graph lower. signature = torch.export.load(str(frozen)).graph_signature @@ -878,6 +908,9 @@ def test_frozen_model_matches_direct_forward( for seed, nloc in ((7, 6), (8, 4)): coord, atype = _batch(seed=seed) coord, atype = coord[:, :nloc], atype[:, :nloc] + if atomic: + rng = np.random.default_rng(seed) + frame["aparam"] = rng.uniform(-1.0, 1.0, (*atype.shape, 1)) expected = model( coord, atype, @@ -902,18 +935,35 @@ def test_frozen_model_matches_direct_forward( ) -def test_change_bias_keeps_the_graph_lower(tmp_path) -> None: - """change-bias rewrites a frozen long-range model on its graph lower.""" +@pytest.mark.parametrize( + ("suffix", "recorded"), [(".pte", "graph"), (".pt2", "dpa4c_canonical")] +) +def test_change_bias_keeps_the_graph_lower( + tmp_path, monkeypatch, suffix: str, recorded: str +) -> None: + """change-bias rewrites a frozen long-range model on its graph lower. + + A .pt2 archive that records the compact canonical lower, a form of CUDA + devices, is rewritten on the lower that the host selects, the graph lower + for this uncompressed model on any device. + """ config = _config("les") model = get_model(config).to(env.DEVICE).eval() - frozen = tmp_path / "long_range.pte" + frozen = tmp_path / f"long_range{suffix}" deserialize_to_file( str(frozen), {"model": model.serialize(), "model_def_script": config}, lower_kind="graph", ) + if recorded != "graph": + read = serialization.serialize_from_file + monkeypatch.setattr( + serialization, + "serialize_from_file", + lambda path: {**read(path), "lower_input_kind": recorded}, + ) bias = [0.5, -1.25] - updated = tmp_path / "long_range_updated.pte" + updated = tmp_path / f"long_range_updated{suffix}" change_bias(str(frozen), bias_value=bias, output=str(updated)) before = DeepPot(str(frozen), auto_batch_size=False) after = DeepPot(str(updated), auto_batch_size=False) @@ -1219,10 +1269,14 @@ def test_pair_graph_skips_padding() -> None: def _periodic_model(kernel: str, **kwargs: object) -> torch.nn.Module: - """A periodic model over the explicit wave numbers of every frame.""" + """A periodic model over the explicit wave numbers of every frame, ``kspace_tolerance`` in ``kwargs``.""" + long_range = { + key: kwargs.pop(key) for key in ("kspace_tolerance",) if key in kwargs + } config = _config(kernel, **kwargs) config["long_range"]["periodic"] = True config["long_range"]["k_grid_mode"] = "ragged" + config["long_range"].update(long_range) return get_model(config).to(env.DEVICE).eval() @@ -1236,18 +1290,23 @@ def _periodic_case() -> tuple[np.ndarray, np.ndarray, np.ndarray]: def _kspace_reference( - coord: np.ndarray, charge: np.ndarray, box: np.ndarray, kernel: str + coord: np.ndarray, + charge: np.ndarray, + box: np.ndarray, + kernel: str, + tolerance: float = 1e-8, ) -> float: """Independent k-space evaluation of the periodic long-range energy, in eV. - E = k_e [ 1/(2V) sum_{k != 0} K_hat(k) |rho(k)|^2 - K(0)/2 sum_i |q_i|^2 ]. + E = k_e [ 1/(2V) sum_{k != 0} K_hat(k) |rho(k)|^2 - K(0)/2 sum_i |q_i|^2 ], + truncated where the weight of the slowest-decaying mode falls below + ``tolerance``, the ``kspace_tolerance`` of the module. """ rec = 2.0 * np.pi * np.linalg.inv(box) # columns g_d: k(n) = rec @ n - # the module's default kspace_tolerance=1e-8 sets the truncation if kernel == "les": - k_max = 2.0 * 0.8 * np.sqrt(np.log(1e8)) + k_max = 2.0 * 0.8 * np.sqrt(np.log(1.0 / tolerance)) else: - k_max = np.sqrt(2.0 * np.log(1e8)) / 0.8 + k_max = np.sqrt(2.0 * np.log(1.0 / tolerance)) / 0.8 n_d = np.ceil(k_max * np.linalg.norm(box, axis=1) / (2 * np.pi)).astype(int) grids = np.meshgrid( *[np.arange(-n, n + 1, dtype=float) for n in n_d], indexing="ij" @@ -1282,16 +1341,28 @@ def _kspace_reference( ) +@pytest.mark.parametrize("tolerance", [None, 1e-4]) @pytest.mark.parametrize("kernel", ["les", "sog"]) -def test_periodic_energy_matches_kspace_reference(kernel: str) -> None: - """With fixed charges the periodic energy is the documented k-space sum.""" - model = _periodic_model(kernel, constraint=None) +def test_periodic_energy_matches_kspace_reference( + kernel: str, tolerance: float | None +) -> None: + """With fixed charges the periodic energy is the documented k-space sum. + + The sum ends at the wave number of ``kspace_tolerance``, the default or a + configured one. + """ + options = {} if tolerance is None else {"kspace_tolerance": tolerance} + model = _periodic_model(kernel, constraint=None, **options) _set_constant_charges(model, [1.0, -0.5]) coord, atype, box = _periodic_case() out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) charge = np.zeros((6, CHANNELS[kernel])) charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] - expected = _kspace_reference(coord[0], charge, box[0], kernel) + expected = _kspace_reference(coord[0], charge, box[0], kernel, tolerance or 1e-8) + if tolerance is not None: + # The configured truncation differs measurably from the default one. + default = _kspace_reference(coord[0], charge, box[0], kernel) + assert abs(expected - default) > 1e-6, (expected, default) np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-9, atol=1e-9) np.testing.assert_allclose( out["atom_energy"].sum().item(), expected, rtol=1e-9, atol=1e-9 diff --git a/source/tests/pt_expt/test_change_bias.py b/source/tests/pt_expt/test_change_bias.py index 1e042b5fd3..f1212324bf 100644 --- a/source/tests/pt_expt/test_change_bias.py +++ b/source/tests/pt_expt/test_change_bias.py @@ -504,6 +504,29 @@ def test_change_bias_pte_preserves_model_def_script(self) -> None: self.assertIn("model_def_script", updated_data) self.assertEqual(updated_data["model_def_script"], original_mds) + def test_change_bias_keeps_the_per_atom_virial(self) -> None: + """change_bias keeps the per-atom virial of a dense-lower archive.""" + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + serialize_from_file, + ) + + data = serialize_from_file(self.shared_pte) + self.assertEqual(data["lower_input_kind"], "nlist") + self.assertFalse(data["do_atomic_virial"]) + with_virial = os.path.join(self.tmpdir, "frozen_atomic_virial.pte") + deserialize_to_file( + with_virial, + {"model": data["model"], "model_def_script": data["model_def_script"]}, + do_atomic_virial=True, + lower_kind="nlist", + ) + output_pte = os.path.join(self.tmpdir, "frozen_atomic_virial_updated.pte") + run_dp(f"dp --pt-expt change-bias {with_virial} -b 0.1 3.2 -o {output_pte}") + updated = serialize_from_file(output_pte) + self.assertEqual(updated["lower_input_kind"], "nlist") + self.assertTrue(updated["do_atomic_virial"]) + def test_change_bias_pt2_preserves_model_def_script(self) -> None: """change_bias on .pt2 should preserve model_def_script (training config).""" from deepmd.pt_expt.utils.serialization import (