Skip to content

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

Description

@wanghan-iapcm

Context

The edge-graph neighbor-list contract (NeighborGraph, deepmd/dpmodel/utils/neighbor_graph/graph.py, per design discussion wanghan-iapcm#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<Canonical=true> 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:

  • PythonNeighborGraph 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 perf(dpa1): optimize graph CUDA inference and deployment #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 perf(dpa1): optimize graph CUDA inference and deployment #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 perf(dpa1): optimize graph CUDA inference and deployment #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#4.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions