deepmd.kernels.triton.sezm.force_assembly

Segmented force / virial assembly from per-edge energy gradients.

Given the per-edge gradient g_e = dE / d(edge_vec_e) of an edge-based energy, the extended force and per-atom virial are

F_k = sum_{dst(e)=k} g_e - sum_{src(e)=k} g_e W_k = 0.5 * sum_{e: k in {src(e), dst(e)}} ( -g_e (x) edge_vec_e ).

The reference assembly issues four index_add scatters (force to both endpoints, half virial to both endpoints) plus a materialized (E, 9) outer product. Row-atomic scatters serialize on the colliding edges of each atom, so this operator performs two CSR segment-reduction launches instead (one over the destination order, one over the source order), each recomputing the per-edge outer product on the fly. One program owns one extended atom; the 12 output scalars (3 force + 9 virial) accumulate in float64 registers over the segment, which both removes the atomic serialization and tightens the summation error over the reference fp32 atomics.

The operator is inference-only in practice: the caller keeps the reference path whenever the force graph must remain differentiable (create_graph), so no autograd formula is registered.

Attributes

FORCE_ASSEMBLY_TRITON_AVAILABLE

Functions

edge_force_assembly(→ tuple[torch.Tensor, torch.Tensor])

Assemble the extended force and per-atom virial from edge gradients.

Module Contents

deepmd.kernels.triton.sezm.force_assembly.FORCE_ASSEMBLY_TRITON_AVAILABLE = True

deepmd.kernels.triton.sezm.force_assembly.edge_force_assembly(g: torch.Tensor, edge_vec: torch.Tensor, dst_order: torch.Tensor, dst_row_ptr: torch.Tensor, src_order: torch.Tensor, src_row_ptr: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

Assemble the extended force and per-atom virial from edge gradients.

Parameters:

gTensor

Per-edge energy gradient dE / d(edge_vec) with shape (E, 3).

edge_vecTensor

Per-edge displacement vectors with shape (E, 3).

dst_order, src_orderTensor

Edge ids sorted by destination / source extended index, each with shape (E,) (from torch.argsort).

dst_row_ptr, src_row_ptrTensor

CSR offsets into the respective order with shape (N_ext + 1,) (from torch.searchsorted on the sorted keys); the length carries the extended-atom count, so the atom axis is never specialized.

Returns:

forceTensor

Extended force with shape (N_ext, 3).

virialTensor

Extended per-atom virial with shape (N_ext, 9), split symmetrically between the two endpoints of each edge.