opt dpa1

Author: OutisLi

deepmd/dpmodel/utils/neighbor_graph/__init__.py

@@ -3,8 +3,9 @@
 
 The unified edge/graph neighbor-list contract and its supporting machinery:
 ``graph`` (the ``NeighborGraph``/``GraphLayout`` contract + derived node-validity
-+ edge padding), ``builder`` (the carry-all ``build_neighbor_graph`` dispatcher +
-the ``from_dense_quartet`` legacy converter), ``segment`` (mask-aware
++ edge padding), ``csr`` (backend-agnostic CSR construction and canonicalization),
+``builder`` (the carry-all ``build_neighbor_graph`` dispatcher + the
+``from_dense_quartet`` legacy converter), ``segment`` (mask-aware
 segment-reduction toolkit), and ``derivatives`` (edge force/virial assembly).
 See the design discussion wanghan-iapcm/deepmd-kit#4.
 """
@@ -24,6 +25,10 @@
     build_neighbor_graph,
     from_dense_quartet,
 )
+from .csr import (
+    build_edge_csr,
+    canonicalize_neighbor_graph,
+)
 from .derivatives import (
     edge_force_virial,
 )
@@ -59,8 +64,10 @@
     "angle_to_node_sum",
     "attach_angles",
     "build_angle_index",
+    "build_edge_csr",
     "build_neighbor_graph",
     "build_neighbor_graph_ase",
+    "canonicalize_neighbor_graph",
     "center_edge_pairs",
     "edge_env_mat",
     "edge_force_virial",

deepmd/dpmodel/utils/neighbor_graph/ase_builder.py

@@ -42,6 +42,9 @@ def build_neighbor_graph_ase(
     box: Array | None,
     rcut: float,
     layout: GraphLayout | None = None,
+    *,
+    with_csr: bool = False,
+    canonicalize: bool = False,
 ) -> NeighborGraph:
     """Build a CARRY-ALL NeighborGraph using ASE's O(N) cell-list search.
 
@@ -66,6 +69,12 @@ def build_neighbor_graph_ase(
         cutoff radius.
     layout
         edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards.
+    with_csr
+        Whether to construct destination/source CSR views for a consumer that
+        requires edge-grouped reductions.
+    canonicalize
+        Whether to reorder every edge field into destination-major form. Implies
+        ``with_csr=True``.
 
     Returns
     -------
@@ -137,5 +146,14 @@ def _to_cpu_numpy(x: Any) -> np.ndarray:
     S_all, nframe_all = S_all[keep], nframe_all[keep]
 
     return neighbor_graph_from_ijs(
-        i_all, j_all, S_all, coord, box, nframe_all, nloc, layout=layout
+        i_all,
+        j_all,
+        S_all,
+        coord,
+        box,
+        nframe_all,
+        nloc,
+        layout=layout,
+        with_csr=with_csr,
+        canonicalize=canonicalize,
     )

deepmd/dpmodel/utils/neighbor_graph/builder.py

@@ -39,6 +39,9 @@
 
 import array_api_compat
 
+from .csr import (
+    build_edge_csr,
+)
 from .graph import (
     GraphLayout,
     NeighborGraph,
@@ -57,6 +60,9 @@ def from_dense_quartet(
     mapping: Array,
     layout: GraphLayout | None = None,
     compact: bool = True,
+    *,
+    with_csr: bool = False,
+    canonicalize: bool = False,
 ) -> NeighborGraph:
     """Convert a legacy extended quartet into a ghost-free NeighborGraph (CONVERTER).
 
@@ -98,6 +104,12 @@ def from_dense_quartet(
         ``src`` pointing at the center (in-range, masked) -- so no ``nonzero`` is
         used and the converter is jit/export-traceable. The masked edges contribute
         zero in a downstream ``segment_sum``, so the descriptor output is unchanged.
+    with_csr
+        Whether to construct destination/source CSR views for a consumer that
+        requires edge-grouped reductions.
+    canonicalize
+        Whether to reorder every edge field into destination-major form. Implies
+        ``with_csr=True``.
 
     Returns
     -------
@@ -108,6 +120,7 @@ def from_dense_quartet(
     """
     if layout is None:
         layout = GraphLayout()
+    with_csr = with_csr or canonicalize
     xp = array_api_compat.array_namespace(extended_coord, nlist, mapping)
     dev = array_api_compat.device(extended_coord)
     nf, nloc, nsel = nlist.shape
@@ -148,11 +161,38 @@ def from_dense_quartet(
         edge_index = xp.astype(xp.stack([src, dst], axis=0), xp.int64)
         edge_mask = valid
         n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev)
+        if not with_csr:
+            return NeighborGraph(
+                n_node=n_node,
+                edge_index=edge_index,
+                edge_vec=edge_vec,
+                edge_mask=edge_mask,
+            )
+        (
+            edge_index,
+            edge_vec,
+            edge_mask,
+            destination_order,
+            destination_row_ptr,
+            source_row_ptr,
+            source_order,
+        ) = build_edge_csr(
+            edge_index,
+            edge_vec,
+            edge_mask,
+            nf * nloc,
+            canonicalize=canonicalize,
+        )
         return NeighborGraph(
             n_node=n_node,
             edge_index=edge_index,
             edge_vec=edge_vec,
             edge_mask=edge_mask,
+            destination_order=destination_order,
+            destination_row_ptr=destination_row_ptr,
+            source_row_ptr=source_row_ptr,
+            source_order=source_order,
+            destination_sorted=canonicalize,
         )
     else:
         # COMPACT: drop invalid slots via nonzero (dynamic shape -> eager only,
@@ -189,11 +229,38 @@ def from_dense_quartet(
             edge_index, edge_vec, layout.edge_capacity, layout.min_edges
         )
         n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev)
+        if not with_csr:
+            return NeighborGraph(
+                n_node=n_node,
+                edge_index=edge_index,
+                edge_vec=edge_vec,
+                edge_mask=edge_mask,
+            )
+        (
+            edge_index,
+            edge_vec,
+            edge_mask,
+            destination_order,
+            destination_row_ptr,
+            source_row_ptr,
+            source_order,
+        ) = build_edge_csr(
+            edge_index,
+            edge_vec,
+            edge_mask,
+            nf * nloc,
+            canonicalize=canonicalize,
+        )
         return NeighborGraph(
             n_node=n_node,
             edge_index=edge_index,
             edge_vec=edge_vec,
             edge_mask=edge_mask,
+            destination_order=destination_order,
+            destination_row_ptr=destination_row_ptr,
+            source_row_ptr=source_row_ptr,
+            source_order=source_order,
+            destination_sorted=canonicalize,
         )
 
 
@@ -203,6 +270,9 @@ def build_neighbor_graph(
     box: Array | None,
     rcut: float,
     layout: GraphLayout | None = None,
+    *,
+    with_csr: bool = False,
+    canonicalize: bool = False,
 ) -> NeighborGraph:
     """Build a CARRY-ALL NeighborGraph DIRECTLY from coordinates (``dense`` search).
 
@@ -236,6 +306,12 @@ def build_neighbor_graph(
         at non-binding ``sel``).
     layout
         edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards.
+    with_csr
+        Whether to construct destination/source CSR views for a consumer that
+        requires edge-grouped reductions.
+    canonicalize
+        Whether to reorder every edge field into destination-major form. Implies
+        ``with_csr=True``.
     """
     from deepmd.dpmodel.utils.nlist import (
         extend_coord_with_ghosts,
@@ -246,6 +322,7 @@ def build_neighbor_graph(
 
     if layout is None:
         layout = GraphLayout()
+    with_csr = with_csr or canonicalize
     xp = array_api_compat.array_namespace(coord, atype)
     dev = array_api_compat.device(coord)
     nf, nloc = atype.shape[:2]
@@ -298,9 +375,36 @@ def build_neighbor_graph(
         edge_index, edge_vec, layout.edge_capacity, layout.min_edges
     )
     n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev)
+    if not with_csr:
+        return NeighborGraph(
+            n_node=n_node,
+            edge_index=edge_index,
+            edge_vec=edge_vec,
+            edge_mask=edge_mask,
+        )
+    (
+        edge_index,
+        edge_vec,
+        edge_mask,
+        destination_order,
+        destination_row_ptr,
+        source_row_ptr,
+        source_order,
+    ) = build_edge_csr(
+        edge_index,
+        edge_vec,
+        edge_mask,
+        nf * nloc,
+        canonicalize=canonicalize,
+    )
     return NeighborGraph(
         n_node=n_node,
         edge_index=edge_index,
         edge_vec=edge_vec,
         edge_mask=edge_mask,
+        destination_order=destination_order,
+        destination_row_ptr=destination_row_ptr,
+        source_row_ptr=source_row_ptr,
+        source_order=source_order,
+        destination_sorted=canonicalize,
     )