Merge branch 'kk_flux' of github.com:stanmoore1/lammps into kk_flux

This commit is contained in:
Stan Moore
2023-10-03 12:07:46 -06:00
2 changed files with 19 additions and 11 deletions

View File

@ -137,13 +137,13 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
int set_flag = 0;
char *str;
if ((str = getenv("SLURM_LOCALID"))) {
if (str = getenv("SLURM_LOCALID")) {
int local_rank = atoi(str);
device = local_rank % ngpus;
if (device >= skip_gpu) device++;
set_flag = 1;
}
if ((str = getenv("MPT_LRANK"))) {
if (str = getenv("FLUX_TASK_LOCAL_ID")) {
if (ngpus > 0) {
int local_rank = atoi(str);
device = local_rank % ngpus;
@ -151,7 +151,7 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
set_flag = 1;
}
}
if ((str = getenv("MV2_COMM_WORLD_LOCAL_RANK"))) {
if (str = getenv("MPT_LRANK")) {
if (ngpus > 0) {
int local_rank = atoi(str);
device = local_rank % ngpus;
@ -159,7 +159,7 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
set_flag = 1;
}
}
if ((str = getenv("OMPI_COMM_WORLD_LOCAL_RANK"))) {
if (str = getenv("MV2_COMM_WORLD_LOCAL_RANK")) {
if (ngpus > 0) {
int local_rank = atoi(str);
device = local_rank % ngpus;
@ -167,7 +167,15 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
set_flag = 1;
}
}
if ((str = getenv("PMI_LOCAL_RANK"))) {
if (str = getenv("OMPI_COMM_WORLD_LOCAL_RANK")) {
if (ngpus > 0) {
int local_rank = atoi(str);
device = local_rank % ngpus;
if (device >= skip_gpu) device++;
set_flag = 1;
}
}
if (str = getenv("PMI_LOCAL_RANK")) {
if (ngpus > 0) {
int local_rank = atoi(str);
device = local_rank % ngpus;

View File

@ -138,9 +138,9 @@ struct PairComputeFunctor {
F_FLOAT fztmp = 0.0;
if (NEIGHFLAG == FULL) {
f(i,0) = 0.0;
f(i,1) = 0.0;
f(i,2) = 0.0;
a_f(i,0) -= f(i,0);
a_f(i,1) -= f(i,1);
a_f(i,2) -= f(i,2);
}
for (int jj = 0; jj < jnum; jj++) {
@ -212,9 +212,9 @@ struct PairComputeFunctor {
F_FLOAT fztmp = 0.0;
if (NEIGHFLAG == FULL) {
f(i,0) = 0.0;
f(i,1) = 0.0;
f(i,2) = 0.0;
a_f(i,0) -= f(i,0);
a_f(i,1) -= f(i,1);
a_f(i,2) -= f(i,2);
}
for (int jj = 0; jj < jnum; jj++) {