Skip to content

Commit 044023e

Browse files
ChiahsinChuChiahsinChu
authored
Merge branch 'devel' into devel-use_aparam_as_mask
2 parents 226c8b9 + 18026eb commit 044023e

File tree

3 files changed

+47
-46
lines changed

3 files changed

+47
-46
lines changed

‎source/api_cc/src/DeepPotPT.cc‎

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,7 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
168168
nlist_data.copy_from_nlist(lmp_list);
169169
nlist_data.shuffle_exclude_empty(fwd_map);
170170
nlist_data.padding();
171-
if (do_message_passing == 1 && nghost > 0) {
171+
if (do_message_passing == 1) {
172172
int nswap = lmp_list.nswap;
173173
torch::Tensor sendproc_tensor =
174174
torch::from_blob(lmp_list.sendproc, {nswap}, int32_option);
@@ -180,10 +180,14 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
180180
torch::from_blob(lmp_list.recvnum, {nswap}, int32_option);
181181
torch::Tensor sendnum_tensor =
182182
torch::from_blob(lmp_list.sendnum, {nswap}, int32_option);
183-
torch::Tensor communicator_tensor = torch::from_blob(
184-
const_cast<void*>(lmp_list.world), {1}, torch::kInt64);
185-
// torch::Tensor communicator_tensor =
186-
// torch::tensor(lmp_list.world, int32_option);
183+
torch::Tensor communicator_tensor;
184+
if (lmp_list.world == 0) {
185+
communicator_tensor = torch::empty({1}, torch::kInt64);
186+
} else {
187+
communicator_tensor = torch::from_blob(
188+
const_cast<void*>(lmp_list.world), {1}, torch::kInt64);
189+
}
190+
187191
torch::Tensor nswap_tensor = torch::tensor(nswap, int32_option);
188192
int total_send =
189193
std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0);
@@ -196,12 +200,6 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
196200
comm_dict.insert("recv_num", recvnum_tensor);
197201
comm_dict.insert("communicator", communicator_tensor);
198202
}
199-
if (do_message_passing == 1 && nghost == 0) {
200-
// for the situation that no ghost atoms (e.g. serial nopbc)
201-
// set the mapping arange(nloc) is enough
202-
auto option = torch::TensorOptions().device(device).dtype(torch::kInt64);
203-
mapping_tensor = at::arange(nloc_real, option).unsqueeze(0);
204-
}
205203
}
206204
at::Tensor firstneigh = createNlistTensor(nlist_data.jlist);
207205
firstneigh_tensor = firstneigh.to(torch::kInt64).to(device);
@@ -224,7 +222,7 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
224222
.to(device);
225223
}
226224
c10::Dict<c10::IValue, c10::IValue> outputs =
227-
(do_message_passing == 1 && nghost > 0)
225+
(do_message_passing == 1)
228226
? module
229227
.run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor,
230228
firstneigh_tensor, mapping_tensor, fparam_tensor,

‎source/lmp/tests/test_lammps_dpa_pt_nopbc.py‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -681,14 +681,6 @@ def test_pair_deepmd_si(lammps_si):
681681
[(["--balance"],), ([],)],
682682
)
683683
def test_pair_deepmd_mpi(balance_args: list):
684-
if balance_args == []:
685-
# python:5331 terminated with signal 11 at PC=7f3e940e3806 SP=7ffd5787edc0. Backtrace:
686-
# /home/runner/work/deepmd-kit/deepmd-kit/dp_test/lib/libdeepmd_op_pt.so(+0x95806)[0x7f3e940e3806]
687-
# /home/runner/work/deepmd-kit/deepmd-kit/dp_test/lib/libdeepmd_op_pt.so(+0x8f76e)[0x7f3e940dd76e]
688-
# /home/runner/work/deepmd-kit/deepmd-kit/dp_test/lib/libdeepmd_op_pt.so(+0x9a38a)[0x7f3e940e838a]
689-
# /home/runner/work/deepmd-kit/deepmd-kit/dp_test/lib/libdeepmd_op_pt.so(_Z9border_opRKN2at6TensorES2_S2_S2_S2_S2_S2_S2_S2_+0x8e)[0x7f3e940dda63]
690-
# /home/runner/work/deepmd-kit/deepmd-kit/dp_test/lib/libdeepmd_op_pt.so(+0xaeac3)[0x7f3e940fcac3]
691-
pytest.skip(reason="Known segfault, see comments for details")
692684
with tempfile.NamedTemporaryFile() as f:
693685
sp.check_call(
694686
[

‎source/op/pt/comm.cc‎

Lines changed: 37 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -87,16 +87,18 @@ class Border : public torch::autograd::Function<Border> {
8787
int mpi_init = 0;
8888
MPI_Initialized(&mpi_init);
8989
int cuda_aware = 1;
90-
int me;
90+
int me = 0;
9191
MPI_Comm world;
9292
int world_size = 0;
93-
unpack_communicator(communicator_tensor, world);
94-
MPI_Comm_rank(world, &me);
95-
MPI_Comm_size(world, &world_size);
93+
if (mpi_init) {
94+
unpack_communicator(communicator_tensor, world);
95+
MPI_Comm_rank(world, &me);
96+
MPI_Comm_size(world, &world_size);
97+
}
9698
MPI_Datatype mpi_type = get_mpi_type<FPTYPE>();
9799
MPI_Request request;
98100
#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
99-
if (world_size != 1) {
101+
if (world_size >= 1) {
100102
int version, subversion;
101103
MPI_Get_version(&version, &subversion);
102104
if (version >= 4) {
@@ -120,11 +122,15 @@ class Border : public torch::autograd::Function<Border> {
120122
for (int iswap = 0; iswap < nswap; ++iswap) {
121123
int nrecv = recvnum[iswap];
122124
int nsend = sendnum[iswap];
123-
torch::Tensor isendlist =
124-
torch::from_blob(sendlist[iswap], {nsend}, int32_options)
125-
.to(recv_g1_tensor.device());
126-
torch::Tensor send_g1_tensor = recv_g1_tensor.index_select(0, isendlist);
127-
FPTYPE* send_g1 = send_g1_tensor.data_ptr<FPTYPE>();
125+
torch::Tensor isendlist;
126+
torch::Tensor send_g1_tensor;
127+
FPTYPE* send_g1;
128+
if (nsend != 0) {
129+
isendlist = torch::from_blob(sendlist[iswap], {nsend}, int32_options)
130+
.to(recv_g1_tensor.device());
131+
send_g1_tensor = recv_g1_tensor.index_select(0, isendlist);
132+
send_g1 = send_g1_tensor.data_ptr<FPTYPE>();
133+
}
128134
#ifdef USE_MPI
129135
if (sendproc[iswap] != me) {
130136
if (nrecv) {
@@ -207,15 +213,17 @@ class Border : public torch::autograd::Function<Border> {
207213
MPI_Initialized(&mpi_init);
208214
int world_size = 0;
209215
int cuda_aware = 1;
216+
int me = 0;
210217
MPI_Comm world;
211-
unpack_communicator(communicator_tensor, world);
212-
int me;
213-
MPI_Comm_rank(world, &me);
214-
MPI_Comm_size(world, &world_size);
218+
if (mpi_init) {
219+
unpack_communicator(communicator_tensor, world);
220+
MPI_Comm_rank(world, &me);
221+
MPI_Comm_size(world, &world_size);
222+
}
215223
MPI_Datatype mpi_type = get_mpi_type<FPTYPE>();
216224
MPI_Request request;
217225
#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
218-
if (world_size != 1) {
226+
if (world_size >= 1) {
219227
int version, subversion;
220228
MPI_Get_version(&version, &subversion);
221229
if (version >= 4) {
@@ -248,17 +256,20 @@ class Border : public torch::autograd::Function<Border> {
248256
int nlocal = nlocal_tensor.item<int>();
249257
int nghost = nghost_tensor.item<int>();
250258
int ntotal = nlocal + nghost;
251-
252-
torch::Tensor send_g1_tensor = d_local_g1_tensor;
253-
254-
int max_recvnum = sendnum_tensor.max().item<int>();
255-
auto options = torch::TensorOptions()
256-
.dtype(d_local_g1_tensor.dtype())
257-
.device(d_local_g1_tensor.device());
258-
torch::Tensor recv_g1_tensor =
259-
torch::empty({max_recvnum, tensor_size}, options);
260-
FPTYPE* recv_g1 = recv_g1_tensor.data_ptr<FPTYPE>();
261-
FPTYPE* send_g1 = send_g1_tensor.data_ptr<FPTYPE>() + ntotal * tensor_size;
259+
torch::Tensor send_g1_tensor;
260+
torch::Tensor recv_g1_tensor;
261+
FPTYPE* recv_g1;
262+
FPTYPE* send_g1;
263+
if (nswap != 0) {
264+
send_g1_tensor = d_local_g1_tensor;
265+
int max_recvnum = sendnum_tensor.max().item<int>();
266+
auto options = torch::TensorOptions()
267+
.dtype(d_local_g1_tensor.dtype())
268+
.device(d_local_g1_tensor.device());
269+
recv_g1_tensor = torch::empty({max_recvnum, tensor_size}, options);
270+
recv_g1 = recv_g1_tensor.data_ptr<FPTYPE>();
271+
send_g1 = send_g1_tensor.data_ptr<FPTYPE>() + ntotal * tensor_size;
272+
}
262273

263274
int end = ntotal;
264275
auto int32_options = torch::TensorOptions().dtype(torch::kInt32);

0 commit comments

Comments
 (0)