Skip to content

SFPG: add a carry-all/reciprocal graph fast path (dst reduction + one broadcast) for the pt_expt graph lower #5957

Description

@wanghan-iapcm

Follow-up from review of #5939 (thread on deepmd/pt/model/descriptor/sezm_nn/edge_cache.py).

Current implementation

The Source Freeze Propagation Gate folds each node's full outgoing-edge set, eta_j = prod_{e: src_e = j} w(r_e). It reduces [log_eta, zero_count] over src. Under domain decomposition a rank only holds edges whose destination it owns, so those per-node partials are rank-incomplete and are completed by a two-phase exchange in _gate_partial_exchange:

  1. border_op_backward — reverse-accumulate ghost rows into their owners;
  2. border_op — forward-broadcast the completed owner values back onto ghosts.

Why source-keyed reduction is required in general

Destination reduction equals source reduction only when the edge set is reciprocal. That does not hold for the sel-capped, nlist-derived edge schema (the PT compatibility path): format_nlist truncates each centre at sel by distance, so a dense-environment atom can drop a neighbour that the sparser neighbour keeps. Measured on an inhomogeneous system (rcut=4.0, sel=6, ghosts mapped to owner images):

sel = 6   edges = 90
NON-RECIPROCAL ordered pairs: 12   [(0,5), (6,2), (3,6), (0,9), (7,2), (3,9)]

So the generic source-keyed semantics must stay for that route.

The fast path this issue tracks

The pt_expt graph-lower route is different, and it is the primary route going forward:

  • it uses the carry-all, sel-free NeighborGraph (every edge inside rcut), so the edge set IS reciprocal;
  • DPA4's dense call() rejects comm_dict before _run_graph, so the communication hook is reached through call_graph, never through a sel-truncated dense graph;
  • therefore the owner holds its complete destination-centred neighbourhood.

On that route the partials can be reduced over dst — already complete for owned nodes — and completed with one forward border_op instead of a reverse-accumulate plus a broadcast. That halves the exchanges, and with them the MPI_Barrier count. It matters more than the payload size suggests: both kernels end in a barrier, and because the exported force graph differentiates through them, their transposes add the matching pair to the force evaluation too.

Requirements

  1. Add it as an explicit, clearly-labelled carry-all/reciprocal fast path. Do not change the generic source-keyed semantics used by the nlist-derived route.
  2. State the reciprocity contract and the routing boundary at the seam (which graphs qualify, and what selects the path), so the assumption is pinned rather than implied.
  3. A negative test that the sel-capped/nlist-derived route still takes the two-phase completion, so the fast path cannot silently capture it.
  4. Parity: fast path vs current two-phase completion, bit-for-bit on a carry-all graph, including a bridging-zone pair straddling a rank boundary (the geometry used in test_dpa4_zbl_parallel.py / test_lammps_dpa4_zbl_pt2.py).
  5. Benchmark across representative system sizes and MPI rank counts. The docs currently make no performance claim about this exchange (feat: SFPG cross-rank completion — capabilities, ZBL bridging and spin+ZBL multi-rank #5939 removed the unsupported "negligible" wording); numbers from this work are what should replace it.

Context

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions