From 06c71e944425c4e6cfe6486f3d24f217f254c44d Mon Sep 17 00:00:00 2001
From: Erik Strand <erik.strand@cba.mit.edu>
Date: Fri, 5 Mar 2021 22:43:06 -0500
Subject: [PATCH] Check SLURM_LOCALID environment variable

This gives you a unique index within a single node. Will be useful for
assigning GPUs to threads.
---
 mpi_info/mpi_info.c | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/mpi_info/mpi_info.c b/mpi_info/mpi_info.c
index 2df42b6..1ec46de 100644
--- a/mpi_info/mpi_info.c
+++ b/mpi_info/mpi_info.c
@@ -1,14 +1,24 @@
 #include "mpi.h"
 #include <stdio.h>
+#include <stdlib.h>
 
 int main(int argc, char *argv[]) {
-    int numtasks, rank, length;
-    char hostname[MPI_MAX_PROCESSOR_NAME];
+    char* local_rank_str = NULL;
+    int local_rank = 0;
+    local_rank_str = getenv("SLURM_LOCALID");
+    if (local_rank_str != NULL) {
+        local_rank = atoi(local_rank_str);
+        printf("slurm local rank = %d\n", local_rank);
+    } else {
+        printf("slurm local rank not defined\n");
+    }
 
+    int n_tasks, rank, length;
+    char host_name[MPI_MAX_PROCESSOR_NAME];
     MPI_Init(&argc, &argv);
-    MPI_Comm_size(MPI_COMM_WORLD, &numtasks);
+    MPI_Comm_size(MPI_COMM_WORLD, &n_tasks);
     MPI_Comm_rank(MPI_COMM_WORLD, &rank);
-    MPI_Get_processor_name(hostname, &length);
-    printf("Number of tasks = %d, rank = %d, running on %s\n", numtasks, rank, hostname);
+    MPI_Get_processor_name(host_name, &length);
+    printf("Number of tasks = %d, rank = %d, running on %s\n", n_tasks, rank, host_name);
     MPI_Finalize();
 }
-- 
GitLab