Skip to content

Commit 58ab3f6

Browse files
LuLu
Lu
authored and
Lu
committed
fix bugs, support cpu environment
1 parent 4961bbf commit 58ab3f6

File tree

5 files changed

+42
-22
lines changed

5 files changed

+42
-22
lines changed

‎source/lib/include/NNPInter.h‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ class NNPInter
127127
bool init_nbor;
128128
std::vector<int> sec_a;
129129
compute_t *array_double;
130+
InternalNeighborList nlist;
130131
NNPAtomMap<VALUETYPE> nnpmap;
131132
unsigned long long *array_longlong;
132133
int *ilist, *jrange, *jlist, *array_int;
@@ -135,7 +136,9 @@ class NNPInter
135136

136137
// function used for neighbor list copy
137138
vector<int> get_sel_a() const;
139+
#ifdef USE_CUDA_TOOLKIT
138140
void update_nbor(const InternalNeighborList & nlist, const int nloc);
141+
#endif
139142
};
140143

141144
class NNPInterModelDevi
@@ -217,8 +220,9 @@ class NNPInterModelDevi
217220

218221
// copy neighbor list info from host
219222
bool init_nbor;
220-
vector<vector<int> > sec;
221223
compute_t *array_double;
224+
vector<vector<int> > sec;
225+
InternalNeighborList nlist;
222226
NNPAtomMap<VALUETYPE> nnpmap;
223227
unsigned long long *array_longlong;
224228
int max_sec_size = 0, max_sec_back = 0;
@@ -229,7 +233,9 @@ class NNPInterModelDevi
229233
void get_max_sec();
230234
vector<vector<int> > get_sel() const;
231235
void cum_sum(const std::vector<std::vector<int32> > n_sel);
236+
#ifdef USE_CUDA_TOOLKIT
232237
void update_nbor(const InternalNeighborList & nlist, const int nloc);
238+
#endif
233239
};
234240

235241

‎source/lib/src/NNPInter.cc‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -683,6 +683,7 @@ NNPInter::~NNPInter() {
683683
#endif
684684
}
685685

686+
#ifdef USE_CUDA_TOOLKIT
686687
void NNPInter::update_nbor(const InternalNeighborList & nlist, const int nloc) {
687688
if (!init_nbor) {
688689
sec_a = cum_sum(get_sel_a());
@@ -742,6 +743,7 @@ void NNPInter::update_nbor(const InternalNeighborList & nlist, const int nloc) {
742743
cudaErrcheck(cudaMemcpy(jrange, &nlist.jrange[0], sizeof(int) * nlist.jrange.size(), cudaMemcpyHostToDevice));
743744
cudaErrcheck(cudaMemcpy(jlist, &nlist.jlist[0], sizeof(int) * nlist.jlist.size(), cudaMemcpyHostToDevice));
744745
}
746+
#endif // USE_CUDA_TOOLKIT
745747

746748
#ifdef USE_CUDA_TOOLKIT
747749
void
@@ -912,7 +914,7 @@ compute (ENERGYTYPE & dener,
912914
{
913915
int nall = dcoord_.size() / 3;
914916
int nloc = nall - nghost;
915-
NNPAtomMap<VALUETYPE> nnpmap (datype_.begin(), datype_.begin() + nloc);
917+
nnpmap = NNPAtomMap<VALUETYPE> (datype_.begin(), datype_.begin() + nloc);
916918
assert (nloc == nnpmap.get_type().size());
917919
validate_fparam_aparam(nloc, fparam, aparam);
918920

@@ -946,7 +948,7 @@ compute (ENERGYTYPE & dener,
946948
nnpmap = NNPAtomMap<VALUETYPE> (datype_.begin(), datype_.begin() + nloc);
947949
assert (nloc == nnpmap.get_type().size());
948950

949-
InternalNeighborList nlist;
951+
// InternalNeighborList nlist;
950952
convert_nlist_lmp_internal (nlist, lmp_list);
951953
shuffle_nlist (nlist, nnpmap);
952954
std::vector<std::pair<string, Tensor>> input_tensors;
@@ -987,7 +989,7 @@ compute (ENERGYTYPE & dener,
987989
const vector<VALUETYPE> & fparam,
988990
const vector<VALUETYPE> & aparam)
989991
{
990-
NNPAtomMap<VALUETYPE> nnpmap (datype_.begin(), datype_.end());
992+
nnpmap = NNPAtomMap<VALUETYPE> (datype_.begin(), datype_.end());
991993
validate_fparam_aparam(nnpmap.get_type().size(), fparam, aparam);
992994

993995
std::vector<std::pair<string, Tensor>> input_tensors;
@@ -1019,10 +1021,10 @@ compute (ENERGYTYPE & dener,
10191021
validate_fparam_aparam(nloc, fparam, aparam);
10201022

10211023
if (ago == 0) {
1022-
NNPAtomMap<VALUETYPE> nnpmap (datype_.begin(), datype_.begin() + nloc);
1024+
nnpmap = NNPAtomMap<VALUETYPE> (datype_.begin(), datype_.begin() + nloc);
10231025
assert (nloc == nnpmap.get_type().size());
10241026

1025-
InternalNeighborList nlist;
1027+
// InternalNeighborList nlist;
10261028
convert_nlist_lmp_internal (nlist, lmp_list);
10271029
shuffle_nlist (nlist, nnpmap);
10281030
std::vector<std::pair<string, Tensor>> input_tensors;
@@ -1240,7 +1242,8 @@ get_max_sec()
12401242
}
12411243
}
12421244

1243-
void
1245+
#ifdef USE_CUDA_TOOLKIT
1246+
void
12441247
NNPInterModelDevi::
12451248
update_nbor(const InternalNeighborList & nlist, const int nloc)
12461249
{
@@ -1303,6 +1306,7 @@ update_nbor(const InternalNeighborList & nlist, const int nloc)
13031306
cudaErrcheck(cudaMemcpy(jrange, &nlist.jrange[0], sizeof(int) * nlist.jrange.size(), cudaMemcpyHostToDevice));
13041307
cudaErrcheck(cudaMemcpy(jlist, &nlist.jlist[0], sizeof(int) * nlist.jlist.size(), cudaMemcpyHostToDevice));
13051308
}
1309+
#endif //USE_CUDA_TOOLKIT
13061310

13071311
void
13081312
NNPInterModelDevi::
@@ -1332,7 +1336,7 @@ compute (ENERGYTYPE & dener,
13321336
{
13331337
if (numb_models == 0) return;
13341338

1335-
NNPAtomMap<VALUETYPE> nnpmap (datype_.begin(), datype_.end());
1339+
nnpmap = NNPAtomMap<VALUETYPE> (datype_.begin(), datype_.end());
13361340
validate_fparam_aparam(nnpmap.get_type().size(), fparam, aparam);
13371341

13381342
std::vector<std::pair<string, Tensor>> input_tensors;
@@ -1389,7 +1393,7 @@ compute (vector<ENERGYTYPE> & all_energy,
13891393
nnpmap = NNPAtomMap<VALUETYPE> (datype_.begin(), datype_.begin() + nloc);
13901394
assert (nloc == nnpmap.get_type().size());
13911395

1392-
InternalNeighborList nlist;
1396+
// InternalNeighborList nlist;
13931397
convert_nlist_lmp_internal (nlist, lmp_list);
13941398
shuffle_nlist (nlist, nnpmap);
13951399
std::vector<std::pair<string, Tensor>> input_tensors;
@@ -1453,7 +1457,7 @@ compute (vector<ENERGYTYPE> & all_energy,
14531457
nnpmap = NNPAtomMap<VALUETYPE> (datype_.begin(), datype_.begin() + nloc);
14541458
assert (nloc == nnpmap.get_type().size());
14551459

1456-
InternalNeighborList nlist;
1460+
// InternalNeighborList nlist;
14571461
convert_nlist_lmp_internal (nlist, lmp_list);
14581462
shuffle_nlist (nlist, nnpmap);
14591463
std::vector<std::pair<string, Tensor>> input_tensors;

‎source/lmp/pair_nnp.cpp‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -288,12 +288,20 @@ void PairNNP::compute(int eflag, int vflag)
288288
#endif
289289
}
290290

291-
int ago = numb_models > 1 ? 0 : neighbor->ago;
292-
291+
// int ago = numb_models > 1 ? 0 : neighbor->ago;
292+
int ago = neighbor->ago;
293+
if (numb_models > 1) {
294+
if (multi_models_no_mod_devi && (out_freq > 0 && update->ntimestep % out_freq == 0)) {
295+
ago = 0;
296+
}
297+
else if (multi_models_mod_devi && (out_freq == 0 || update->ntimestep % out_freq != 0)) {
298+
ago = 0;
299+
}
300+
}
293301
// compute
294-
bool single_model = (numb_models == 1);
295-
bool multi_models_no_mod_devi = (numb_models > 1 && (out_freq == 0 || update->ntimestep % out_freq != 0));
296-
bool multi_models_mod_devi = (numb_models > 1 && (out_freq > 0 && update->ntimestep % out_freq == 0));
302+
single_model = (numb_models == 1);
303+
multi_models_no_mod_devi = (numb_models > 1 && (out_freq == 0 || update->ntimestep % out_freq != 0));
304+
multi_models_mod_devi = (numb_models > 1 && (out_freq > 0 && update->ntimestep % out_freq == 0));
297305
if (do_ghost) {
298306
LammpsNeighborList lmp_list (list->inum, list->ilist, list->numneigh, list->firstneigh);
299307
if (single_model || multi_models_no_mod_devi) {

‎source/lmp/pair_nnp.h.in‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,9 @@ private:
8282
int dim_aparam;
8383
int out_each;
8484
int out_rel;
85+
bool single_model;
86+
bool multi_models_mod_devi;
87+
bool multi_models_no_mod_devi;
8588
#ifdef HIGH_PREC
8689
vector<double > fparam;
8790
vector<double > aparam;

‎source/op/CMakeLists.txt‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,15 @@ file(GLOB OP_PY *.py)
1111
if (BUILD_CPP_IF)
1212
if (USE_CUDA_TOOLKIT)
1313
add_library(${LIB_DEEPMD_OP} SHARED ${OP_CUDA_SRC})
14+
add_subdirectory(cuda)
15+
find_package(CUDA REQUIRED)
16+
include_directories(${CUDA_INCLUDE_DIRS})
17+
set (EXTRA_LIBS ${EXTRA_LIBS} deepmd_op_cuda)
18+
target_link_libraries (${LIB_DEEPMD_OP} ${EXTRA_LIBS})
19+
target_link_libraries (${LIB_DEEPMD_OP} ${CUDA_LIBRARIES})
1420
else (USE_CUDA_TOOLKIT)
1521
add_library(${LIB_DEEPMD_OP} SHARED ${OP_SRC})
1622
endif (USE_CUDA_TOOLKIT)
17-
18-
add_subdirectory(cuda)
19-
find_package(CUDA REQUIRED)
20-
include_directories(${CUDA_INCLUDE_DIRS})
21-
set (EXTRA_LIBS ${EXTRA_LIBS} deepmd_op_cuda)
22-
target_link_libraries (${LIB_DEEPMD_OP} ${EXTRA_LIBS})
23-
target_link_libraries (${LIB_DEEPMD_OP} ${CUDA_LIBRARIES})
2423
endif (BUILD_CPP_IF)
2524

2625
if (BUILD_PY_IF)

0 commit comments

Comments
 (0)