From 884b9dd9bd855acd55b83a75dfd95a3cf4eaf0f8 Mon Sep 17 00:00:00 2001
From: Chris Knight <knight@x4402c1s0b0n0.hostmgmt.cm.aurora.alcf.anl.gov>
Date: Fri, 19 Jul 2024 15:34:02 +0000
Subject: [PATCH] add PALS env var

---
 src/KOKKOS/kokkos.cpp | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/KOKKOS/kokkos.cpp b/src/KOKKOS/kokkos.cpp
index 58b9436af6..1c36169002 100644
--- 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 "