ICODE-ADC/lammps
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support for HPE MPI to Kokkos package

Browse Source
This commit is contained in:
Stan Moore
2019-09-09 10:14:30 -06:00
parent cfd72aacce
commit 51786a459f
3 changed files with 21 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/KOKKOS/kokkos.cpp
View File

@ -113,23 +113,37 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
}
iarg += 2;
int set_flag = 0;
char *str;
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"))) {
int local_rank = atoi(str);
device = local_rank % ngpus;
if (device >= skip_gpu) device++;
set_flag = 1;
}
if ((str = getenv("MV2_COMM_WORLD_LOCAL_RANK"))) {
int local_rank = atoi(str);
device = local_rank % ngpus;
if (device >= skip_gpu) device++;
set_flag = 1;
}
if ((str = getenv("OMPI_COMM_WORLD_LOCAL_RANK"))) {
int local_rank = atoi(str);
device = local_rank % ngpus;
if (device >= skip_gpu) device++;
set_flag = 1;
}
if (ngpus > 1 && !set_flag)
error->all(FLERR,"Could not determine local MPI rank for multiple "
"GPUs with Kokkos CUDA because MPI library not recognized");
} else if (strcmp(arg[iarg],"t") == 0 ||
strcmp(arg[iarg],"threads") == 0) {
nthreads = atoi(arg[iarg+1]);