Skip to content

Commit 255ac9b

Browse files
OutisLiOutisLi
committed
fix cuda
1 parent 324c4ee commit 255ac9b

2 files changed

Lines changed: 111 additions & 96 deletions

File tree

‎deepmd/pt/utils/nv_nlist.py‎

Lines changed: 80 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@
1818
annotations,
1919
)
2020

21+
import contextlib
2122
import logging
2223
from typing import (
24+
TYPE_CHECKING,
2325
Any,
2426
)
2527

@@ -37,6 +39,11 @@
3739

3840
log = logging.getLogger(__name__)
3941

42+
if TYPE_CHECKING:
43+
from collections.abc import (
44+
Iterator,
45+
)
46+
4047

4148
def is_nv_available() -> bool:
4249
"""Whether the ``nvalchemiops`` Toolkit-Ops neighbor list is importable."""
@@ -69,6 +76,17 @@ def choose_nv_nlist_method(nloc: int, *, periodic: bool = True) -> str:
6976
return "batch_naive"
7077

7178

79+
@contextlib.contextmanager
80+
def _input_device_context(device: torch.device) -> Iterator[None]:
81+
"""Run third-party kernels with both default and current devices pinned."""
82+
if device.type == "cuda":
83+
with torch.device(device), torch.cuda.device(device):
84+
yield
85+
else:
86+
with torch.device(device):
87+
yield
88+
89+
7290
class NvNeighborList(NeighborList):
7391
"""Neighbor-list strategy using the ``nvalchemiops`` kernels.
7492
@@ -95,71 +113,72 @@ def build(
95113
neighbor_list,
96114
)
97115

98-
nf, nloc = atype.shape[:2]
99116
device = coord.device
100-
target_neighbors = int(sum(sel))
101-
search_capacity = target_neighbors
102-
total_atoms = nf * nloc
103-
coord = coord.reshape(nf, nloc, 3)
104-
periodic = box is not None
105-
if not periodic:
106-
cell = None
107-
pbc = None
108-
else:
109-
cell = box.reshape(nf, 3, 3).to(device=device, dtype=coord.dtype)
110-
coord = normalize_coord(coord, cell)
111-
pbc = torch.ones((nf, 3), dtype=torch.bool, device=device)
112-
positions_for_nlist = coord.reshape(total_atoms, 3).detach()
113-
batch_idx = torch.arange(
114-
nf, dtype=torch.int32, device=device
115-
).repeat_interleave(nloc)
116-
batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc
117-
method = choose_nv_nlist_method(nloc, periodic=periodic)
118-
119-
# Grow the search capacity until all neighbors fit so the distance-sort
120-
# below selects the true nearest ``sum(sel)``.
121-
while True:
122-
nlist_result = neighbor_list(
123-
positions_for_nlist,
124-
float(rcut),
117+
with _input_device_context(device):
118+
nf, nloc = atype.shape[:2]
119+
target_neighbors = int(sum(sel))
120+
search_capacity = target_neighbors
121+
total_atoms = nf * nloc
122+
coord = coord.reshape(nf, nloc, 3)
123+
periodic = box is not None
124+
if not periodic:
125+
cell = None
126+
pbc = None
127+
else:
128+
cell = box.reshape(nf, 3, 3).to(device=device, dtype=coord.dtype)
129+
coord = normalize_coord(coord, cell)
130+
pbc = torch.ones((nf, 3), dtype=torch.bool, device=device)
131+
positions_for_nlist = coord.reshape(total_atoms, 3).detach()
132+
batch_idx = torch.arange(
133+
nf, dtype=torch.int32, device=device
134+
).repeat_interleave(nloc)
135+
batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc
136+
method = choose_nv_nlist_method(nloc, periodic=periodic)
137+
138+
# Grow the search capacity until all neighbors fit so the distance-sort
139+
# below selects the true nearest ``sum(sel)``.
140+
while True:
141+
nlist_result = neighbor_list(
142+
positions_for_nlist,
143+
float(rcut),
144+
cell=cell,
145+
pbc=pbc,
146+
batch_idx=batch_idx,
147+
batch_ptr=batch_ptr,
148+
method=method,
149+
max_neighbors=int(search_capacity),
150+
return_neighbor_list=False,
151+
wrap_positions=False,
152+
)
153+
if len(nlist_result) == 2:
154+
neighbor_matrix, num_neighbors = nlist_result
155+
shifts = torch.zeros(
156+
(*neighbor_matrix.shape, 3),
157+
dtype=torch.int32,
158+
device=device,
159+
)
160+
else:
161+
neighbor_matrix, num_neighbors, shifts = nlist_result
162+
max_found = (
163+
int(num_neighbors.max().item()) if num_neighbors.numel() > 0 else 0
164+
)
165+
if max_found <= search_capacity:
166+
break
167+
search_capacity = max(max_found, _grow_search_capacity(search_capacity))
168+
169+
extended_coord, extended_atype, mapping, nlist = _matrix_to_extended_inputs(
170+
coord=coord,
171+
atype=atype,
125172
cell=cell,
126-
pbc=pbc,
127-
batch_idx=batch_idx,
128-
batch_ptr=batch_ptr,
129-
method=method,
130-
max_neighbors=int(search_capacity),
131-
return_neighbor_list=False,
132-
wrap_positions=False,
173+
nloc=nloc,
174+
neighbor_matrix=neighbor_matrix,
175+
num_neighbors=num_neighbors,
176+
shifts=shifts,
133177
)
134-
if len(nlist_result) == 2:
135-
neighbor_matrix, num_neighbors = nlist_result
136-
shifts = torch.zeros(
137-
(*neighbor_matrix.shape, 3),
138-
dtype=torch.int32,
139-
device=device,
140-
)
141-
else:
142-
neighbor_matrix, num_neighbors, shifts = nlist_result
143-
max_found = (
144-
int(num_neighbors.max().item()) if num_neighbors.numel() > 0 else 0
178+
nlist = _truncate_to_sel_compiled(
179+
extended_coord, nlist, target_neighbors, float(rcut)
145180
)
146-
if max_found <= search_capacity:
147-
break
148-
search_capacity = max(max_found, _grow_search_capacity(search_capacity))
149-
150-
extended_coord, extended_atype, mapping, nlist = _matrix_to_extended_inputs(
151-
coord=coord,
152-
atype=atype,
153-
cell=cell,
154-
nloc=nloc,
155-
neighbor_matrix=neighbor_matrix,
156-
num_neighbors=num_neighbors,
157-
shifts=shifts,
158-
)
159-
nlist = _truncate_to_sel_compiled(
160-
extended_coord, nlist, target_neighbors, float(rcut)
161-
)
162-
return extended_coord, extended_atype, nlist, mapping
181+
return extended_coord, extended_atype, nlist, mapping
163182

164183

165184
def _grow_search_capacity(capacity: int) -> int:

‎source/tests/pt/model/test_nv_nlist.py‎

Lines changed: 31 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
builder at the nlist level (edge topology + geometry).
99
"""
1010

11-
import contextlib
1211
import unittest
1312
from unittest.mock import (
1413
patch,
@@ -24,6 +23,7 @@
2423
)
2524
from deepmd.pt.utils.nv_nlist import (
2625
NvNeighborList,
26+
_input_device_context,
2727
)
2828

2929
_NV_AVAILABLE = nv_nlist.is_nv_available()
@@ -166,24 +166,18 @@ def _assert_nv_matches_native(
166166
sel: list[int],
167167
force_cell_list: bool = False,
168168
) -> None:
169-
# native: (extended_coord, extended_atype, mapping, nlist)
170-
native = extend_input_and_build_neighbor_list(
171-
coord,
172-
atype,
173-
rcut,
174-
sel,
175-
mixed_types=True,
176-
box=box,
177-
)
178-
# NeighborList strategy: (extended_coord, extended_atype, nlist, mapping)
179-
builder = NvNeighborList()
180-
# Pin the current CUDA device so the Toolkit-Ops backend launches there.
181-
device_ctx = (
182-
torch.cuda.device(coord.device)
183-
if coord.is_cuda
184-
else contextlib.nullcontext()
185-
)
186-
with device_ctx:
169+
with _input_device_context(coord.device):
170+
# native: (extended_coord, extended_atype, mapping, nlist)
171+
native = extend_input_and_build_neighbor_list(
172+
coord,
173+
atype,
174+
rcut,
175+
sel,
176+
mixed_types=True,
177+
box=box,
178+
)
179+
# NeighborList strategy: (extended_coord, extended_atype, nlist, mapping)
180+
builder = NvNeighborList()
187181
if force_cell_list:
188182
with (
189183
patch.object(nv_nlist, "NV_CELL_LIST_THRESHOLD", 1),
@@ -192,23 +186,25 @@ def _assert_nv_matches_native(
192186
nv = builder.build(coord, atype, box, rcut, sel)
193187
else:
194188
nv = builder.build(coord, atype, box, rcut, sel)
195-
native_coord, _, native_mapping, native_nlist = native
196-
nv_coord, nv_atype, nv_nlist_out, nv_mapping = nv
197-
# The strategy trims to sum(sel) itself, so the width is fixed.
198-
self.assertEqual(nv_nlist_out.shape[-1], sum(sel))
199-
self.assertTrue(
200-
torch.equal(
201-
_edge_topology_from_extended(native_mapping, native_nlist),
202-
_edge_topology_from_extended(nv_mapping, nv_nlist_out),
189+
native_coord, _, native_mapping, native_nlist = native
190+
nv_coord, nv_atype, nv_nlist_out, nv_mapping = nv
191+
# The strategy trims to sum(sel) itself, so the width is fixed.
192+
self.assertEqual(nv_nlist_out.shape[-1], sum(sel))
193+
self.assertTrue(
194+
torch.equal(
195+
_edge_topology_from_extended(native_mapping, native_nlist),
196+
_edge_topology_from_extended(nv_mapping, nv_nlist_out),
197+
)
203198
)
204-
)
205-
torch.testing.assert_close(
206-
_edge_geometry_from_extended(native_coord, native_mapping, native_nlist),
207-
_edge_geometry_from_extended(nv_coord, nv_mapping, nv_nlist_out),
208-
atol=1.0e-10,
209-
rtol=1.0e-10,
210-
)
211-
_assert_extended_atype_matches_mapping(self, atype, nv_atype, nv_mapping)
199+
torch.testing.assert_close(
200+
_edge_geometry_from_extended(
201+
native_coord, native_mapping, native_nlist
202+
),
203+
_edge_geometry_from_extended(nv_coord, nv_mapping, nv_nlist_out),
204+
atol=1.0e-10,
205+
rtol=1.0e-10,
206+
)
207+
_assert_extended_atype_matches_mapping(self, atype, nv_atype, nv_mapping)
212208

213209
def test_cell_list_matches_native(self) -> None:
214210
"""The ``batch_cell_list`` method (forced via the threshold) matches the

0 commit comments

Comments
 (0)