root/opal/mca/pmix/cray/pmix_cray_pmap_parser.c

DEFINITIONS

1. find_my_node
2. find_lrs
3. pmix_cray_parse_pmap
4. dump_lrs
5. main

   1 /* -*- Mode: C; c-basic-offset:4 ; -*- */
   2 /*
   3  *
   4  * Copyright (c) 2013      Mellanox Technologies, Inc.
   5  *                         All rights reserved.
   6  * Copyright (c) 2014      Intel, Inc. All rights reserved.
   7  * Copyright (c) 2016 Cisco Systems, Inc.  All rights reserved.
   8  * Copyright (c) 2016      Los Alamos National Security, LLC. All rights
   9  *                         reserved.
  10  * $COPYRIGHT$
  11  * Additional copyrights may follow
  12  *
  13  * $HEADER$
  14  *
  15  */
  16 
  17 #include "pmix_cray_pmap_parser.h"
  18 
  19 /**
  20  pmi2 process mapping is returned as a
  21  comma separated list of tuples:
  22 ex: (vector,(0,4,4),(0,4,1))
  23 slurm cyclic distro of 4 ranks over 2 nodes:
  24 (vector,(0,2,1),(0,2,1))
  25 slurm block distro of 4 ranks over 2 nodes:
  26 (vector,(0,2,2))
  27 
  28 
  29  Format of each tuple is (base, H, L), where
  30  H is number of nodes spawned by tuple,
  31  L is number of ranks per node,
  32  base is offset from node 0.
  33 
  34  Tuple can be visualized as a rectangle on two
  35  dimensional (Hosts, Local Ranks) plane:
  36 
  37            ------------------------------------ Hosts ->
  38            |              H
  39            |           +--------+
  40            |<- base -->|        |
  41            |           |        | L
  42            |           +--------+
  43         Local Ranks
  44            V
  45 
  46 Note that ranks increase by column. Tuple (0,2,3) looks like:
  47 0 3
  48 1 4
  49 2 5
  50 
  51 */
  52 #include <stdio.h>
  53 #include <stdlib.h>
  54 #include <string.h>
  55 static int find_my_node(char *map, int me)
  56 {
  57     int abs_rank;
  58     int base, H, L;
  59     char *p;
  60 
  61     p = map;
  62     abs_rank = 0;
  63     while (NULL != (p = strstr(p+1, ",("))) {
  64         if (3 != sscanf(p, ",(%d,%d,%d)", &base, &H, &L)) {
  65             return -1;
  66         }
  67         if (me >= abs_rank && me < abs_rank + H*L) {
  68             /* found my rectangle, compute node */
  69             return base + (me - abs_rank)/L;
  70         }
  71         abs_rank += H*L;
  72     }
  73     return -1;
  74 }
  75 
  76 static int *find_lrs(char *map, int my_node, int *nlrs)
  77 {
  78     int abs_rank;
  79     int base, H, L;
  80     char *p;
  81     int *lrs;
  82     int max_lr;
  83     int i;
  84 
  85     p = map;
  86     abs_rank = 0;
  87     *nlrs = 0;
  88     max_lr = 16;
  89     lrs = malloc(max_lr * sizeof(int));
  90     while (NULL != (p = strstr(p+1, ",("))) {
  91         if (3 != sscanf(p, ",(%d,%d,%d)", &base, &H, &L)) {
  92             free(lrs);
  93             return NULL;
  94         }
  95         if (base <= my_node && my_node < base + H) {
  96             if (*nlrs + L >= max_lr) {
  97                 lrs = realloc(lrs, (max_lr + L) * sizeof(int));
  98                 if (NULL == lrs) {
  99                     *nlrs = 0;
 100                     free(lrs);
 101                     return NULL;
 102                 }
 103                 max_lr += L;
 104             }
 105             /* skip (my_node - base) columns of L elems,
 106              * numbers in my column are local to me
 107              */
 108             for (i = 0; i < L; i++) {
 109                 lrs[*nlrs] = (my_node - base) * L + i + abs_rank;
 110                 (*nlrs) ++;
 111             }
 112         }
 113         abs_rank += H*L;
 114     }
 115 
 116     if (0 == *nlrs) {
 117         free(lrs);
 118         lrs = 0;
 119     }
 120     return lrs;
 121 }
 122 
 123 /**
 124  * @param pmap process map as returned by PMI_process_mapping
 125  *             attribute
 126  * @param my_rank
 127  * @param node set to my node id
 128  * @param nlrs set to the number of local ranks returned
 129  *
 130  * @return array that contains ranks local to my_rank or NULL
 131  * on failure. Array must be freed by the caller.
 132  */
 133 int *pmix_cray_parse_pmap(char *pmap, int my_rank,
 134                           int *node, int *nlrs)
 135 {
 136     char *p;
 137 
 138     p = strstr(pmap, "(vector");
 139     if (NULL == p) {
 140         return NULL;
 141     }
 142 
 143     *node = find_my_node(p, my_rank);
 144     if (0 > *node) {
 145        return NULL;
 146     }
 147 
 148     return find_lrs(p, *node, nlrs);
 149 }
 150 
 151 
 152 #ifdef STANDALONE_TEST
 153 #include <assert.h>
 154 static void dump_lrs(int *lrs, int me, int node, int n)
 155 {
 156     int i;
 157 
 158     printf("Total %d ranks/node, node %d me %d\n", n, node, me);
 159     for (i = 0; i < n; i++) {
 160         printf("%d ", lrs[i]);
 161     }
 162     printf("\n");
 163     free(lrs);
 164 }
 165 
 166 int main(int argc, char **argv)
 167 {
 168     int me, n, node;
 169     int *lrs;
 170     char *pmap;
 171     int a1[] = {0, 1};
 172     int a2[] = {2, 3};
 173     int a3[] = {0, 2};
 174     int a4[] = {1, 3};
 175     int a5[] = {0,1,3,2,16,17};
 176     int a6[] = {8,9,10,11,19};
 177 
 178 
 179     if (argc == 3) {
 180         me = atoi(argv[1]);
 181         lrs = orte_grpcomm_pmi2_parse_pmap(argv[2], me, &node, &n);
 182         if (NULL == lrs) {
 183             printf("can not parse pmap\n");
 184             exit(1);
 185         }
 186         dump_lrs(lrs, me, node, n);
 187         exit(0);
 188     }
 189     /* built in cases */
 190 
 191     pmap = "(vector,(0,2,2))";
 192     me = 1;
 193     lrs = orte_grpcomm_pmi2_parse_pmap(pmap, me, &node, &n);
 194     assert(lrs);
 195     assert(n == 2);
 196     assert(memcmp(lrs, a1, 2) == 0);
 197     free(lrs);
 198 
 199 
 200     pmap = "(vector,(0,2,2))";
 201     me = 2;
 202     lrs = orte_grpcomm_pmi2_parse_pmap(pmap, me, &node, &n);
 203     assert(lrs);
 204     assert(n == 2);
 205     assert(memcmp(lrs, a2, 2) == 0);
 206     free(lrs);
 207 
 208 
 209     /* cyclic distro which skips node 0 */
 210     pmap = "(vector,(1,2,1),(1,2,1))";
 211     me = 0;
 212     lrs = orte_grpcomm_pmi2_parse_pmap(pmap, me, &node, &n);
 213     assert(lrs);
 214     assert(n == 2);
 215     assert(memcmp(lrs, a3, n) == 0);
 216     free(lrs);
 217 
 218     pmap = "(vector,(1,2,1),(1,2,1))";
 219     me = 3;
 220     lrs = orte_grpcomm_pmi2_parse_pmap(pmap, me, &node, &n);
 221     assert(lrs);
 222     assert(n == 2);
 223     assert(memcmp(lrs, a4, n) == 0);
 224     free(lrs);
 225 
 226     pmap = "(vector,(0,4,4),(0,1,2),(1,3,1))";
 227     me = 3;
 228     lrs = orte_grpcomm_pmi2_parse_pmap(pmap, me, &node, &n);
 229     assert(lrs);
 230     assert(n == 6);
 231     assert(memcmp(lrs, a5, n) == 0);
 232     free(lrs);
 233 
 234     pmap = "(vector,(0,4,4),(0,1,2),(1,3,1))";
 235     me = 10;
 236     lrs = orte_grpcomm_pmi2_parse_pmap(pmap, me, &node, &n);
 237     assert(lrs);
 238     assert(n == 5);
 239     assert(memcmp(lrs, a6, n) == 0);
 240     free(lrs);
 241     return 0;
 242 }
 243 #endif