deepmd.pt_expt.kernels.dpa4c

deepmd.pt_expt.kernels.dpa4c#

Fused compressed operators for the DPA4C graph lower.

The bindings are device neutral: the CUDA kernels live in source/op/pt/dpa4c/*.cu and the CPU kernels in source/op/pt/dpa4c/*_cpu.cc, and the dispatcher selects between them.

Submodules#

Functions#

`canonical_model_eligible`(→ bool)	Return whether a model can use the compact source-only graph ABI.
`dpa4c_canonical_compress_energy_force`(...)	Evaluate compressed DPA4C from a compact canonical edge stream.
`build_radial_table`(→ tuple[torch.Tensor, torch.Tensor])	Tabulate the composed DPA4C radial branch and its mode profiles.
`dpa4c_graph_compress`(→ torch.Tensor)	Evaluate the compressed DPA4C graph descriptor.
`dpa4c_graph_compress_energy_force`(...)	Evaluate compressed DPA4C energy, force, virial and magnetic force.
`ef_op_available`(→ bool)	Return whether the descriptor, fitting, and force operators are present.
`ensure_registered`(→ None)	Register fake and autograd implementations when the op is available.
`mega_eligible`(→ bool)	Return whether the descriptor has a compiled fp32 specialization.
`op_available`(→ bool)	Return whether the backend device carries the DPA4C compressed operator.

Package Contents#

deepmd.pt_expt.kernels.dpa4c.canonical_model_eligible(model: Any) → bool[source]#

Return whether a model can use the compact source-only graph ABI.

The compact ABI exists so that a device-resident neighbor list can reach the descriptor without a host round trip, which is a property of the CUDA / Kokkos deployment. Only that backend implements its operators, so a CPU target keeps the generic graph lower, whose operators it does implement and whose extra cost is one index tensor.

deepmd.pt_expt.kernels.dpa4c.dpa4c_canonical_compress_energy_force(descriptor: Any, fitting: Any, graph: deepmd.pt_expt.utils.canonical_graph.CanonicalGraph, atype: torch.Tensor, ownership: torch.Tensor, atom_bias: torch.Tensor, do_atomic_virial: bool, spin: torch.Tensor | None = None, pair_table: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor][source]#

Evaluate compressed DPA4C from a compact canonical edge stream.

The compact ABI carries only source indices and CSR row pointers, so the descriptor operator addresses edges by position and omits every mask, padding-type, and table-tail check that the cutoff-compacted graph already guarantees. The backward reuses the forward state in place.

Parameters:

descriptor: Compressed pt_expt DPA4C descriptor.
fitting: Eligible pt_expt energy fitting network.
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,).
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.
spin: Per-node magnetic moments with shape (N, 3) for a spin-conditioned descriptor, or None.
pair_table: Series constants of an analytical pair potential with shape ((T + 1) ** 2, 8), or None. The fused operator adds its energy to the node energies and its slope to the edge gradient.

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.
force_mag: Per-node magnetic force with shape (N, 3), or an empty tensor. The on-site part closes inside the fused operator; the neighbour part is emitted per edge and reduced onto source nodes here, where the source CSR is in scope.

Raises:

ValueError: If the model or the compiled operators do not support the compact path.

deepmd.pt_expt.kernels.dpa4c.build_radial_table(descriptor: Any, stride: float = 0.002) → tuple[torch.Tensor, torch.Tensor][source]#

Tabulate the composed DPA4C radial branch and its mode profiles.

The table width is channels + radial_modes: the leading block holds the shared radial map read by every degree, and the trailing block holds the shared mode profiles that each ordered type pair mixes.

Parameters:

descriptor: pt_expt DPA4C descriptor whose DPA4 radial modules define the table.
stride: Uniform distance spacing in Å.

Returns:

table: Quintic coefficients with shape (S, 6 * (channels + radial_modes)), fp32.
info: CPU metadata [stride, table_max, rcut, eps, degree_floor], fp64.

Raises:

ValueError: If stride is not positive.

deepmd.pt_expt.kernels.dpa4c.dpa4c_graph_compress(descriptor: Any, graph: Any, atype: torch.Tensor, spin: torch.Tensor | None = None) → torch.Tensor[source]#

Evaluate the compressed DPA4C graph descriptor.

Parameters:

descriptor: Compressed pt_expt DPA4C descriptor.
graph: NeighborGraph with destination CSR topology.
atype: Flat node types with shape (N,).
spin: Per-node magnetic moments with shape (N, 3), or None.

Returns:

torch.Tensor: Degree-wise invariant descriptor with shape (N, descriptor.get_dim_out()), fp32.

Raises:

ValueError: If the graph lacks destination CSR topology.

deepmd.pt_expt.kernels.dpa4c.dpa4c_graph_compress_energy_force(descriptor: Any, fitting: Any, graph: Any, atype: torch.Tensor, ownership: torch.Tensor, atom_bias: torch.Tensor, node_capacity: int, do_atomic_virial: bool, spin: torch.Tensor | None = None, pair_table: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor][source]#

Evaluate compressed DPA4C energy, force, virial and magnetic force.

Parameters:

descriptor: Compressed pt_expt DPA4C descriptor.
fitting: Eligible pt_expt energy fitting network.
graph: NeighborGraph with destination and source CSR topology.
atype: Flat node atom types with shape (N,).
ownership: Boolean mask selecting energy-contributing nodes with shape (N,).
atom_bias: Combined atomic energy bias with shape (ntypes,).
node_capacity: Force-scatter node capacity.
do_atomic_virial: Whether to return per-node virials.
spin: Per-node magnetic moments with shape (N, 3), or None.
pair_table: Series constants of an analytical pair potential with shape ((T + 1) ** 2, 8), or None. The backward scan adds its energy to the node energies and its slope to the edge gradient.

Returns:

energy: Per-frame energy with shape (F, 1), fp64.
atom_energy: Per-node energy with shape (N, 1), fp64.
force: Per-node force with shape (N, 3), fp32.
virial: Per-frame virial with shape (F, 3, 3), fp32.
atom_virial: Per-node virial with shape (N, 3, 3) or an empty tensor.
force_mag: Per-node magnetic force with shape (N, 3), or an empty tensor.

Raises:

ValueError: If the graph lacks destination or source CSR topology.

deepmd.pt_expt.kernels.dpa4c.ef_op_available(spin: bool = False) → bool[source]#

Return whether the descriptor, fitting, and force operators are present.

Parameters:

spin: Whether the caller needs the native-spin variant.

Returns:

bool: Whether every operator of the fused energy-force route is registered.

deepmd.pt_expt.kernels.dpa4c.ensure_registered() → None[source]#: Register fake and autograd implementations when the op is available.

deepmd.pt_expt.kernels.dpa4c.mega_eligible(descriptor: Any) → bool[source]#

Return whether the descriptor has a compiled fp32 specialization.

Native spin is a compiled variant of the same kernels rather than a separate operator, and every spin width follows from the channel count, so a spin-conditioned descriptor is eligible on exactly the conditions a spin-free one is. Each condition below is a width the compiled operator specializes on.

Eligibility is a property of the descriptor alone: it decides whether the compression artifacts are worth building, which a snapshot does once for every device that may later consume it. Whether a device has the kernel to consume them is op_available.

Parameters:

descriptor: Descriptor instance, which need not be a DPA4C.

Returns:

bool: Whether the compiled operator covers this configuration.

deepmd.pt_expt.kernels.dpa4c.op_available(spin: bool = False) → bool[source]#

Return whether the backend device carries the DPA4C compressed operator.

Native spin is a compiled variant of the same operator rather than a separate one, so it needs no separate availability probe – except on the CPU, whose kernels carry no magnetic branch: those families need a source-major counterpart of the destination scan, and no CPU deployment asks for them yet. A spin-conditioned descriptor therefore keeps the reference path there.

Parameters:

spin: Whether the caller needs the native-spin variant.

Returns:

bool: Whether the operator is registered for the backend device.