one more env var for local rank

2022-08-09 17:35:22 +00:00
parent fb2dd26d45
commit d5170a1703
1 changed files with 6 additions and 0 deletions
--- a/src/KOKKOS/kokkos.cpp
+++ b/src/KOKKOS/kokkos.cpp
@ -160,6 +160,12 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
        if (device >= skip_gpu) device++;
        set_flag = 1;
      }
+      if ((str = getenv("PMI_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 "