deepmodeling / deepmodeling/deepmd-kit

refactor: enforce canonical-graph invariants via factory-only construction (NeighborGraph / GraphTensorPack)

Open
#5,777 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

## Context

The edge-graph neighbor-list contract (`NeighborGraph`, `deepmd/dpmodel/utils/neighbor_graph/graph.py`, per design discussion wanghan-iapcm/deepmd-kit#4) supports an optional destination-major *canonical* normal form: the edge payload is physically sorted by destination, masked edges move to the suffix, `destination_order` collapses to the identity, and `destination_sorted=True`. This normal form exists to feed the compressed-DPA1 CUDA kernel, whose exported `edge_at_csr_position` indexes edges by CSR position directly (skips a gather) and therefore **requires** destination-major input.

Introduced/extended in #5758.

## Problem: the canonical invariant is enforced by convention, not by the type

Across both backends, the coupling

```
destination_sorted == True ==> destination_order is identity AND edge payload is physically destination-major
```

holds only because the single code path that sets it (`attach_edge_csr(canonicalize=True)` / `canonicalizeGraphPayload`) happens to set all fields consistently. Nothing in the types prevents an inconsistent "fake-canonical" graph:

- **Python** — `NeighborGraph` is a plain dataclass; a caller can pass `destination_sorted=True` together with a non-identity `destination_order`, and the object is accepted.
- **C++** — `GraphTensorPack` (`source/api_cc/include/commonPT.h`) is a public-field aggregate mutated field-by-field, and does not even carry a `destination_sorted` flag: "canonical" is the purely implicit convention that `destination_order` happens to be `arange`. The LAMMPS-Kokkos device path (`DeepPotPTExpt.cc`, the non-canonicalizing `build_graph_csr` branch) relies on the Kokkos producer emitting destination-major edges, with nothing asserting it (see the non-blocking robustness note in #5758). A future change to the Kokkos edge layout, or reuse by another non-sorted edge producer, would silently compute the descriptor from mismatched edges.

There is also a second, parallel compact representation for deployment — `DPA1CanonicalGraph` (`deepmd/pt_expt/utils/canonical_graph.py`) and its C++ twin `CanonicalGraphTensorPack` — which re-encodes the same topology in a different field set. These are today produced only by `compactCanonicalGraph()` / the canonicalize path in practice, but that is not enforced by construction either.

Related review threads on #5758: the C++ robustness gap and the dataclass field-order awkwardness are both symptoms of this "too much loosely-structured optional state, enforced by convention" pressure.

## Proposal: factory-only canonical construction

Make an inconsistent canonical graph unrepresentable by routing all canonical construction through a single blessed factory, so the raw constructor cannot reach the canonical state.

### Python (`dpmodel` + `pt_expt`)
- Either (a) keep `NeighborGraph` but remove `destination_sorted` / CSR views from the public init and expose a classmethod (`NeighborGraph.canonical(graph, n_nodes)`) as the only setter of the canonical state; or (b) introduce a distinct `CanonicalNeighborGraph` type constructible only via `from_graph(...)`, so functions that require canonical input declare it in their signature.
- Make `DPA1CanonicalGraph` a derived-view-only type (constructible solely via a `from_neighbor_graph()` factory), never independently assembled.

### C++ (`source/api_cc`)
- Cheap first step (also closes the #5758 robustness gap): add a `bool destination_sorted` member to `GraphTensorPack`, set it only inside `canonicalizeGraphPayload` / `buildGraphCSR(..., destination_sorted=true)`, and `TORCH_CHECK(pack.destination_sorted)` at the compressed-kernel entry so the LAMMPS-Kokkos assumption fails loudly instead of miscomputing.
- Durable step: convert `CanonicalGraphTensorPack` from a public-field aggregate into a class with a private constructor + static factory (`CanonicalGraphTensorPack::compact(...)`) + const accessors, and make the compressed-kernel entry / `run_model_graph` take the canonical type by reference so the type in the signature proves destination-major. (torch tensors are refcounted handles, so accessor wrapping is zero-cost.)

## Scope / tasks
- [ ] Python: factory-only canonical construction for `NeighborGraph` (option a or b)
- [ ] Python: `DPA1CanonicalGraph` as derived-view-only
- [ ] C++: `destination_sorted` flag + `TORCH_CHECK` at kernel entry (closes the #5758 robustness gap)
- [ ] C++: `CanonicalGraphTensorPack` private-ctor + static factory + accessors; typed compressed-kernel entry
- [ ] Tests covering both branches (canonical vs non-canonical construction), including rejection of an inconsistent canonical graph

## Invariant to preserve
`deepmd/dpmodel/**` must never import torch-only kernel/graph code (currently clean).

Refs: #5758, design discussion wanghan-iapcm/deepmd-kit#4.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.