add PALS env var

2024-07-19 15:34:02 +00:00
parent bae822fea7
commit 884b9dd9bd
1 changed files with 8 additions and 0 deletions
--- a/src/KOKKOS/kokkos.cpp
+++ b/src/KOKKOS/kokkos.cpp
@ -149,6 +149,14 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp)
          set_flag = 1;
        }
      }
+      if ((str = getenv("PALS_LOCAL_RANKID"))) {
+        if (ngpus > 0) {
+          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 "