Function DP_DPA4ComputeCanonicalGraphGPU

Function DP_DPA4ComputeCanonicalGraphGPU#

Function Documentation#

void DP_DPA4ComputeCanonicalGraphGPU(DP_DeepPot *dp, double *d_atom_energy, double *d_force, double *d_virial, double *d_atom_virial, const int64_t *d_atype, const void *d_source, DP_CanonicalIndexWidth index_width, const float *d_edge_vec, const int64_t *d_destination_row_ptr, const int64_t *d_source_row_ptr, const int64_t *d_source_order, const double *d_coord, void *world, int64_t nloc, int64_t nall_nodes, int64_t edge_storage, void *stream)#

Evaluate a compact canonical graph on the model device: the device-resident inference of the DPA4 family, whose compressed DPA4C archives are frozen with the compact canonical lower, as are compressed DPA1 archives.

Since

API version 31

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.

Parameters:

  • dp – [in] The DP to use.

  • d_atom_energy – [out] Per-local-atom energy, shape (nloc).

  • d_force – [out] Per-node force, shape (nall_nodes, 3).

  • d_virial – [out] Virial of the graph, nine components, or NULL when it is not requested.

  • d_atom_virial – [out] Per-node virial, shape (nall_nodes, 9), or NULL when the per-atom virial is not requested.

  • d_atype – [in] Per-node atom types, shape (nall_nodes).

  • d_source – [in] Source node per edge storage slot.

  • index_width – [in] Width of each source index, four or eight bytes.

  • d_edge_vec – [in] FP32 edge vectors, shape (edge_storage, 3).

  • d_destination_row_ptr – [in] Destination CSR offsets.

  • d_source_row_ptr – [in] Source CSR offsets. May be NULL together with d_source_order and d_atom_virial unless DP_DeepPotCanonicalGraphRequiresSourceCSR reports otherwise.

  • d_source_order – [in] Source-grouped int64 edge storage positions.

  • d_coord – [in] Node positions in the length unit of the model, shape (nall_nodes, 3), in the order of d_atype, or NULL. Only a model with a long-range module (DP_DeepPotHasLongRange) reads them, and it requires them; every other model ignores them.

  • world – [in] 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.

  • nloc – [in] Number of owned local nodes.

  • nall_nodes – [in] Total local-plus-halo node count.

  • edge_storage – [in] Number of edge storage slots.

  • stream – [in] 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. 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.