Template Function deepmd::assembleGraph

Template Function deepmd::assembleGraph#

Function Documentation#

template<typename VALUETYPE>
inline GraphTensorPack deepmd::assembleGraph(const SkinTopology &topology, const VALUETYPE *coord, const std::int64_t *atype, const int *keep_table, const int ntypes, const double rcut, const bool edge_vec_fp32, const bool with_source_csr, const torch::Device &device, GraphAssemblyScratch &scratch)#

Assemble the destination-major neighbor graph for one geometry.

Replaces a chain of roughly twenty tensor operations – gather, difference, norm, comparison, nonzero, three index_select, three cat, a cast and three sorts – with one threaded pass over the cached topology. The chain moved about half a gigabyte per step on a production system and several of its stages were single-threaded.

Two passes over the cached topology cost less than one pass that stages its survivors: the coordinates a host-sized neighbor list addresses stay resident in L2, so gathering a candidate twice is cheaper than writing it to a staging buffer and copying it back. Measured on an 8000-atom step, the two passes together take 0.56 ms against 2.09 ms for the staged form.

The cutoff filter and the model-level type exclusion are one predicate, so an excluded edge is never allocated rather than allocated and masked. That is what the sort-based path achieved by moving masked edges into a suffix outside every row, and it is exactly equivalent for a consumer that reads the row pointers.

Two masked edges terminate the payload so that the exported graph never observes an empty edge axis.

Parameters:

  • topology – Cached skin topology, destination-grouped.

  • coord – Extended coordinates, length 3 * nall, in the index space of topology.source_ext.

  • atype – Node types, length topology.node_count; read only when a keep table is supplied.

  • keep_table – Flat (ntypes + 1)^2 type-pair keep table, or null.

  • ntypes – Number of real atom types.

  • rcut – Model cutoff.

  • edge_vec_fp32 – Whether the payload carries float32 displacements.

  • with_source_csr – Whether the source-major permutation is needed.

  • device – Target device for the returned tensors.

  • scratch – Reused count and offset storage.

Returns:

The graph pack, with an identity destination permutation.