10f4d7579c722bd601781c29e397a696ed71f02f
[charm.git] / src / conv-core / cpuaffinity.c
1
2 /*
3  This scheme relies on using IP address to identify nodes and assigning 
4 cpu affinity.  
5
6  when CMK_NO_SOCKETS, which is typically on cray xt3 and bluegene/L.
7  There is no hostname for the compute nodes.
8  *
9  * last updated 3/20/2010   Gengbin Zheng
10  * new options +pemap +commmap takes complex pattern of a list of cores
11 */
12
13 #define _GNU_SOURCE
14
15 #include "converse.h"
16 #include "sockRoutines.h"
17
18 #define DEBUGP(x)   /* CmiPrintf x;  */
19
20 #if CMK_HAS_SETAFFINITY || defined (_WIN32) || CMK_HAS_BINDPROCESSOR
21
22 #include <stdlib.h>
23 #include <stdio.h>
24 #include <unistd.h>
25
26 #ifdef _WIN32
27 #include <windows.h>
28 #include <winbase.h>
29 #else
30 #define _GNU_SOURCE
31 #include <sched.h>
32 //long sched_setaffinity(pid_t pid, unsigned int len, unsigned long *user_mask_ptr);
33 //long sched_getaffinity(pid_t pid, unsigned int len, unsigned long *user_mask_ptr);
34 #endif
35
36 #if defined(__APPLE__) 
37 #include <Carbon/Carbon.h> /* Carbon APIs for Multiprocessing */
38 #endif
39
40 #if defined(ARCH_HPUX11) ||  defined(ARCH_HPUX10)
41 #include <sys/mpctl.h>
42 #endif
43
44
45 #define MAX_EXCLUDE      16
46 static int excludecore[MAX_EXCLUDE] = {-1};
47 static int excludecount = 0;
48
49 static int affinity_doneflag = 0;
50
51 static int in_exclude(int core)
52 {
53   int i;
54   for (i=0; i<excludecount; i++) if (core == excludecore[i]) return 1;
55   return 0;
56 }
57
58 static void add_exclude(int core)
59 {
60   if (in_exclude(core)) return;
61   CmiAssert(excludecount < MAX_EXCLUDE);
62   excludecore[excludecount++] = core;
63 }
64
65 #if CMK_HAS_BINDPROCESSOR
66 #include <sys/processor.h>
67 #endif
68
69 #define SET_MASK(cpuid)    \
70   /* set the affinity mask if possible */      \
71   if ((cpuid / 8) > len) {      \
72     printf("Mask size too small to handle requested CPU ID\n");   \
73     return -1;      \
74   } else {    \
75     mask = 1 << cpuid;   /* set the affinity mask exclusively to one CPU */ \
76   }
77
78
79 /* This implementation assumes the default x86 CPU mask size used by Linux */
80 /* For a large SMP machine, this code should be changed to use a variable sized   */
81 /* CPU affinity mask buffer instead, as the present code will fail beyond 32 CPUs */
82 int set_cpu_affinity(unsigned int cpuid) {
83   unsigned long mask = 0xffffffff;
84   unsigned int len = sizeof(mask);
85   int retValue = 0;
86   int pid;
87
88  #ifdef _WIN32
89    HANDLE hProcess;
90  #endif
91  
92 #ifdef _WIN32
93   SET_MASK(cpuid)
94   hProcess = GetCurrentProcess();
95   if (SetProcessAffinityMask(hProcess, mask) == 0) {
96     return -1;
97   }
98 #elif CMK_HAS_BINDPROCESSOR
99   pid = getpid();
100   if (bindprocessor(BINDPROCESS, pid, cpuid) == -1) return -1;
101 #else
102   SET_MASK(cpuid)
103
104   /* PID 0 refers to the current process */
105   if (sched_setaffinity(0, len, &mask) < 0) {
106     perror("sched_setaffinity");
107     return -1;
108   }
109 #endif
110
111   return 0;
112 }
113
114 #if CMK_SMP
115 int set_thread_affinity(int cpuid) {
116   unsigned long mask = 0xffffffff;
117   unsigned int len = sizeof(mask);
118
119 #ifdef _WIN32
120   HANDLE hThread;
121 #endif  
122   
123 #ifdef _WIN32
124   SET_MASK(cpuid)
125   hThread = GetCurrentThread();
126   if (SetThreadAffinityMask(hThread, mask) == 0) {
127     return -1;
128   }
129 #elif  CMK_HAS_PTHREAD_SETAFFINITY
130   int s, j;
131   cpu_set_t cpuset;
132   pthread_t thread;
133
134   thread = pthread_self();
135
136   CPU_ZERO(&cpuset);
137   CPU_SET(cpuid, &cpuset);
138
139   s = pthread_setaffinity_np(thread, sizeof(cpu_set_t), &cpuset);
140   if (s != 0) {
141     perror("pthread_setaffinity");
142     return -1;
143   }
144 #elif CMK_HAS_BINDPROCESSOR
145   if (bindprocessor(BINDTHREAD, thread_self(), cpuid) != 0)
146     return -1;
147 #else
148   return set_cpu_affinity(cpuid);
149 #endif
150
151   return 0;
152 }
153 #endif
154
155 int CmiSetCPUAffinity(int mycore)
156 {
157   int core = mycore;
158   if (core < 0) {
159     core = CmiNumCores() + core;
160   }
161   if (core < 0) {
162     CmiError("Error: Invalid cpu affinity core number: %d\n", mycore);
163     CmiAbort("CmiSetCPUAffinity failed");
164   }
165   /* set cpu affinity */
166 #if CMK_SMP
167   return set_thread_affinity(core);
168 #else
169   return set_cpu_affinity(core);
170   /* print_cpu_affinity(); */
171 #endif
172 }
173
174 /* This implementation assumes the default x86 CPU mask size used by Linux */
175 /* For a large SMP machine, this code should be changed to use a variable sized   */
176 /* CPU affinity mask buffer instead, as the present code will fail beyond 32 CPUs */
177 int print_cpu_affinity() {
178 #ifdef _WIN32
179   unsigned long pMask, sMask;
180   HANDLE hProcess = GetCurrentProcess();
181   if(GetProcessAffinityMask(hProcess, &pMask, &sMask)){
182         perror("On Windows: GetProcessAffinityMask");
183     return -1;
184   }
185   
186  CmiPrintf("[%d] CPU affinity mask is: 0x%08lx\n", CmiMyPe(), pMask);
187   
188 #elif CMK_HAS_BINDPROCESSOR
189   printf("[%d] CPU affinity mask is unknown for AIX. \n", CmiMyPe());
190 #else
191   unsigned long mask;
192   unsigned int len = sizeof(mask);
193
194   /* PID 0 refers to the current process */
195   if (sched_getaffinity(0, len, &mask) < 0) {
196     perror("sched_getaffinity");
197     return -1;
198   }
199
200   CmiPrintf("[%d] CPU affinity mask is: 0x%08lx\n", CmiMyPe(), mask);
201 #endif
202   return 0;
203 }
204
205 #if CMK_SMP
206 int print_thread_affinity() {
207   unsigned long mask;
208   size_t len = sizeof(mask);
209
210 #if  CMK_HAS_PTHREAD_SETAFFINITY
211   int s, j;
212   cpu_set_t cpuset;
213   pthread_t thread;
214   char str[256], pe[16];
215
216   thread = pthread_self();
217   s = pthread_getaffinity_np(thread, sizeof(cpu_set_t), &cpuset);
218   if (s != 0) {
219     perror("pthread_getaffinity");
220     return -1;
221   }
222
223   sprintf(str, "[%d] %s affinity is: ", CmiMyPe(), CmiMyPe()>=CmiNumPes()?"communication pthread":"pthread");
224   for (j = 0; j < CPU_SETSIZE; j++)
225         if (CPU_ISSET(j, &cpuset)) {
226             sprintf(pe, " %d", j);
227             strcat(str, pe);
228         }
229   CmiPrintf("%s\n", str);
230 #endif
231   return 0;
232 }
233 #endif
234
235 int CmiPrintCPUAffinity()
236 {
237 #if CMK_SMP
238   return print_thread_affinity();
239 #else
240   return print_cpu_affinity();
241 #endif
242 }
243
244   /* returns which core is running on, only works on Linux */
245 int CmiOnCore(void) {
246   FILE *fp;
247   int core=-1;
248   char str[128];
249   int n;
250
251   fp = fopen("/proc/self/stat", "r");
252   if (fp == NULL) return -1;
253   for (n=0; n<39; n++)  {
254     fscanf(fp, "%s", str);
255   }
256   fclose(fp);
257   core = atoi(str);
258
259   return core;
260 }
261
262 static int cpuAffinityHandlerIdx;
263 static int cpuAffinityRecvHandlerIdx;
264
265 typedef struct _hostnameMsg {
266   char core[CmiMsgHeaderSizeBytes];
267   int pe;
268   skt_ip_t ip;
269   int ncores;
270   int rank;
271   int seq;
272 } hostnameMsg;
273
274 typedef struct _rankMsg {
275   char core[CmiMsgHeaderSizeBytes];
276   int *ranks;                  /* PE => core rank mapping */
277   int *nodes;                  /* PE => node number mapping */
278 } rankMsg;
279
280 static rankMsg *rankmsg = NULL;
281 static CmmTable hostTable;
282 static CmiNodeLock affLock = NULL;
283
284 /* called on PE 0 */
285 static void cpuAffinityHandler(void *m)
286 {
287   static int count = 0;
288   static int nodecount = 0;
289   hostnameMsg *rec;
290   hostnameMsg *msg = (hostnameMsg *)m;
291   hostnameMsg *tmpm;
292   int tag, tag1, pe, myrank;
293   int npes = CmiNumPes();
294
295 /*   for debug
296   char str[128];
297   skt_print_ip(str, msg->ip);
298   printf("hostname: %d %s\n", msg->pe, str);
299 */
300   CmiAssert(CmiMyPe()==0 && rankmsg != NULL);
301   tag = *(int*)&msg->ip;
302   pe = msg->pe;
303   if ((rec = (hostnameMsg *)CmmProbe(hostTable, 1, &tag, &tag1)) != NULL) {
304     CmiFree(msg);
305   }
306   else {
307     rec = msg;
308     rec->seq = nodecount;
309     nodecount++;                          /* a new node record */
310     CmmPut(hostTable, 1, &tag, msg);
311   }
312   myrank = rec->rank%rec->ncores;
313   while (in_exclude(myrank)) {             /* skip excluded core */
314     myrank = (myrank+1)%rec->ncores;
315     rec->rank ++;
316   }
317   rankmsg->ranks[pe] = myrank;             /* core rank */
318   rankmsg->nodes[pe] = rec->seq;           /* on which node */
319   rec->rank ++;
320   count ++;
321   if (count == CmiNumPes()) {
322     /* CmiPrintf("Cpuaffinity> %d unique compute nodes detected! \n", CmmEntries(hostTable)); */
323     tag = CmmWildCard;
324     while (tmpm = CmmGet(hostTable, 1, &tag, &tag1)) CmiFree(tmpm);
325     CmmFree(hostTable);
326 #if 1
327     /* bubble sort ranks on each node according to the PE number */
328     {
329     int i,j;
330     for (i=0; i<npes-1; i++)
331       for(j=i+1; j<npes; j++) {
332         if (rankmsg->nodes[i] == rankmsg->nodes[j] && 
333               rankmsg->ranks[i] > rankmsg->ranks[j]) 
334         {
335           int tmp = rankmsg->ranks[i];
336           rankmsg->ranks[i] = rankmsg->ranks[j];
337           rankmsg->ranks[j] = tmp;
338         }
339       }
340     }
341 #endif
342     CmiSyncBroadcastAllAndFree(sizeof(rankMsg)+CmiNumPes()*sizeof(int)*2, (void *)rankmsg);
343   }
344 }
345
346 /* called on each processor */
347 static void cpuAffinityRecvHandler(void *msg)
348 {
349   int myrank, mynode;
350   rankMsg *m = (rankMsg *)msg;
351   m->ranks = (int *)((char*)m + sizeof(rankMsg));
352   m->nodes = (int *)((char*)m + sizeof(rankMsg) + CmiNumPes()*sizeof(int));
353   myrank = m->ranks[CmiMyPe()];
354   mynode = m->nodes[CmiMyPe()];
355
356   /*CmiPrintf("[%d %d] set to core #: %d\n", CmiMyNode(), CmiMyPe(), myrank);*/
357
358   if (-1 != CmiSetCPUAffinity(myrank)) {
359     DEBUGP(("Processor %d is bound to core #%d on node #%d\n", CmiMyPe(), myrank, mynode));
360   }
361   else
362     CmiPrintf("Processor %d set affinity failed!\n", CmiMyPe());
363
364   CmiFree(m);
365 }
366
367 #if defined(_WIN32) && ! defined(__CYGWIN__)
368   /* strtok is thread safe in VC++ */
369 #define strtok_r(x,y,z) strtok(x,y)
370 #endif
371
372 static int search_pemap(char *pecoremap, int pe)
373 {
374   int *map = (int *)malloc(CmiNumPes()*sizeof(int));
375   char *ptr = NULL;
376   int i, count;
377   char *str;
378
379   char *mapstr = (char*)malloc(strlen(pecoremap)+1);
380   strcpy(mapstr, pecoremap);
381
382   str = strtok_r(mapstr, ",", &ptr);
383   count = 0;
384   while (str)
385   {
386       int hasdash=0, hascolon=0;
387       int start, end, stride=1;
388       for (i=0; i<strlen(str); i++) {
389           if (str[i] == '-' && i!=0) hasdash=1;
390           if (str[i] == ':') hascolon=1;
391       }
392       if (hasdash) {
393           if (hascolon) {
394             if (sscanf(str, "%d-%d:%d", &start, &end, &stride) != 3)
395                  printf("Warning: Check the format of \"%s\".\n", str);
396           }
397           else {
398             if (sscanf(str, "%d-%d", &start, &end) != 2)
399                  printf("Warning: Check the format of \"%s\".\n", str);
400           }
401       }
402       else {
403           sscanf(str, "%d", &start);
404           end = start;
405       }
406       for (i = start; i<=end; i+=stride) {
407         map[count++] = i;
408         if (count == CmiNumPes()) break;
409       }
410       if (count == CmiNumPes()) break;
411       str = strtok_r(NULL, ",", &ptr);
412   }
413   i = map[pe % count];
414
415   free(map);
416   free(mapstr);
417   return i;
418 }
419
420 #if CMK_CRAYXT
421 extern int getXTNodeID(int mype, int numpes);
422 #endif
423
424 void CmiInitCPUAffinity(char **argv)
425 {
426   static skt_ip_t myip;
427   int ret, i, exclude;
428   hostnameMsg  *msg;
429   char *pemap = NULL;
430   char *commap = NULL;
431   char *coremap = NULL;
432  
433   int show_affinity_flag;
434   int affinity_flag = CmiGetArgFlagDesc(argv,"+setcpuaffinity",
435                                                 "set cpu affinity");
436
437   while (CmiGetArgIntDesc(argv,"+excludecore", &exclude, "avoid core when setting cpuaffinity")) 
438     if (CmiMyRank() == 0) add_exclude(exclude);
439   
440     /* obsolete */
441   CmiGetArgStringDesc(argv, "+coremap", &coremap, "define core mapping");
442   if (coremap!=NULL && excludecount>0)
443     CmiAbort("Charm++> +excludecore and +coremap can not be used togetehr!");
444
445   CmiGetArgStringDesc(argv, "+pemap", &pemap, "define pe to core mapping");
446   if (pemap!=NULL && (coremap != NULL || excludecount>0))
447     CmiAbort("Charm++> +pemap can not be used with either +excludecore or +coremap.");
448   CmiGetArgStringDesc(argv, "+commap", &commap, "define comm threads to core mapping");
449   show_affinity_flag = CmiGetArgFlagDesc(argv,"+showcpuaffinity",
450                                                 "print cpu affinity");
451
452   cpuAffinityHandlerIdx =
453        CmiRegisterHandler((CmiHandler)cpuAffinityHandler);
454   cpuAffinityRecvHandlerIdx =
455        CmiRegisterHandler((CmiHandler)cpuAffinityRecvHandler);
456
457   if (CmiMyRank() ==0) {
458      affLock = CmiCreateLock();
459   }
460
461   if (!affinity_flag) {
462     if (show_affinity_flag) CmiPrintCPUAffinity();
463     return;
464   }
465   else if (CmiMyPe() == 0) {
466      CmiPrintf("Charm++> cpu affinity enabled. \n");
467      if (excludecount > 0) {
468        CmiPrintf("Charm++> cpuaffinity excludes core: %d", excludecore[0]);
469        for (i=1; i<excludecount; i++) CmiPrintf(" %d", excludecore[i]);
470        CmiPrintf(".\n");
471      }
472      if (coremap!=NULL)
473        CmiPrintf("Charm++> cpuaffinity core map : %s\n", coremap);
474      if (pemap!=NULL)
475        CmiPrintf("Charm++> cpuaffinity PE-core map : %s\n", pemap);
476   }
477
478   if (CmiMyPe() >= CmiNumPes()) {         /* this is comm thread */
479       /* comm thread either can float around, or pin down to the last rank.
480          however it seems to be reportedly slower if it is floating */
481     CmiNodeAllBarrier();
482     if (commap != NULL) {
483       int mycore = search_pemap(commap, CmiMyPe()-CmiNumPes());
484       printf("Charm++> set comm %d on node %d to core #%d\n", CmiMyPe()-CmiNumPes(), CmiMyNode(), mycore); 
485       if (-1 == CmiSetCPUAffinity(mycore))
486         CmiAbort("set_cpu_affinity abort!");
487     }
488     else {
489     /* if (CmiSetCPUAffinity(CmiNumCores()-1) == -1) CmiAbort("set_cpu_affinity abort!"); */
490     }
491     if (coremap == NULL && pemap == NULL) {
492 #if CMK_MACHINE_PROGRESS_DEFINED
493     while (affinity_doneflag < CmiMyNodeSize())  CmiNetworkProgress();
494 #else
495 #if CMK_SMP
496     #error "Machine progress call needs to be implemented for cpu affinity!"
497 #endif
498 #endif
499     }
500     CmiNodeAllBarrier();
501     if (show_affinity_flag) CmiPrintCPUAffinity();
502     return;    /* comm thread return */
503   }
504
505   if (pemap != NULL) {
506     int mycore = search_pemap(pemap, CmiMyPe());
507     printf("Charm++> set PE %d on node %d to core #%d\n", CmiMyPe(), CmiMyNode(), mycore); 
508     if (mycore >= CmiNumCores()) {
509       CmiPrintf("Error> Invalid core number %d, only have %d cores (0-%d) on the node. \n", mycore, CmiNumCores(), CmiNumCores()-1);
510       CmiAbort("Invalid core number");
511     }
512     if (CmiSetCPUAffinity(mycore) == -1) CmiAbort("set_cpu_affinity abort!");
513     CmiNodeAllBarrier();
514     CmiNodeAllBarrier();
515     if (show_affinity_flag) CmiPrintCPUAffinity();
516     return;
517   }
518
519   if (coremap!= NULL) {
520     /* each processor finds its mapping */
521     int i, ct=1, myrank;
522     for (i=0; i<strlen(coremap); i++) if (coremap[i]==',' && i<strlen(coremap)-1) ct++;
523 #if CMK_SMP
524     ct = CmiMyRank()%ct;
525 #else
526     ct = CmiMyPe()%ct;
527 #endif
528     i=0;
529     while(ct>0) if (coremap[i++]==',') ct--;
530     myrank = atoi(coremap+i);
531     printf("Charm++> set PE%d on node #%d to core #%d\n", CmiMyPe(), CmiMyNode(), myrank); 
532     if (myrank >= CmiNumCores()) {
533       CmiPrintf("Error> Invalid core number %d, only have %d cores (0-%d) on the node. \n", myrank, CmiNumCores(), CmiNumCores()-1);
534       CmiAbort("Invalid core number");
535     }
536     if (CmiSetCPUAffinity(myrank) == -1) CmiAbort("set_cpu_affinity abort!");
537     CmiNodeAllBarrier();
538     CmiNodeAllBarrier();
539     return;
540   }
541
542     /* get my ip address */
543   if (CmiMyRank() == 0)
544   {
545 #if CMK_CRAYXT
546     ret = getXTNodeID(CmiMyPe(), CmiNumPes());
547     memcpy(&myip, &ret, sizeof(int));
548 #elif CMK_HAS_GETHOSTNAME
549     myip = skt_my_ip();        /* not thread safe, so only calls on rank 0 */
550 #else
551     CmiAbort("Can not get unique name for the compute nodes. \n");
552 #endif
553   }
554   CmiNodeAllBarrier();
555
556     /* prepare a msg to send */
557   msg = (hostnameMsg *)CmiAlloc(sizeof(hostnameMsg));
558   CmiSetHandler((char *)msg, cpuAffinityHandlerIdx);
559   msg->pe = CmiMyPe();
560   msg->ip = myip;
561   msg->ncores = CmiNumCores();
562   DEBUGP(("PE %d's node has %d number of cores. \n", CmiMyPe(), msg->ncores));
563   msg->rank = 0;
564   CmiSyncSendAndFree(0, sizeof(hostnameMsg), (void *)msg);
565
566   if (CmiMyPe() == 0) {
567     int i;
568     hostTable = CmmNew();
569     rankmsg = (rankMsg *)CmiAlloc(sizeof(rankMsg)+CmiNumPes()*sizeof(int)*2);
570     CmiSetHandler((char *)rankmsg, cpuAffinityRecvHandlerIdx);
571     rankmsg->ranks = (int *)((char*)rankmsg + sizeof(rankMsg));
572     rankmsg->nodes = (int *)((char*)rankmsg + sizeof(rankMsg) + CmiNumPes()*sizeof(int));
573     for (i=0; i<CmiNumPes(); i++) {
574       rankmsg->ranks[i] = 0;
575       rankmsg->nodes[i] = -1;
576     }
577
578     for (i=0; i<CmiNumPes(); i++) CmiDeliverSpecificMsg(cpuAffinityHandlerIdx);
579   }
580
581     /* receive broadcast from PE 0 */
582   CmiDeliverSpecificMsg(cpuAffinityRecvHandlerIdx);
583   CmiLock(affLock);
584   affinity_doneflag++;
585   CmiUnlock(affLock);
586   CmiNodeAllBarrier();
587
588   if (show_affinity_flag) CmiPrintCPUAffinity();
589 }
590
591 #else           /* not supporting affinity */
592
593
594 int CmiPrintCPUAffinity()
595 {
596   CmiPrintf("Warning: CmiPrintCPUAffinity not supported.\n");
597 }
598
599 void CmiInitCPUAffinity(char **argv)
600 {
601   char *pemap = NULL;
602   char *commap = NULL;
603   char *coremap = NULL;
604   int excludecore = -1;
605   int affinity_flag = CmiGetArgFlagDesc(argv,"+setcpuaffinity",
606                                                 "set cpu affinity");
607   while (CmiGetArgIntDesc(argv,"+excludecore",&excludecore, "avoid core when setting cpuaffinity"));
608   CmiGetArgStringDesc(argv, "+coremap", &coremap, "define core mapping");
609   CmiGetArgStringDesc(argv, "+pemap", &pemap, "define pe to core mapping");
610   CmiGetArgStringDesc(argv, "+commap", &commap, "define comm threads to core mapping");
611   if (affinity_flag && CmiMyPe()==0)
612     CmiPrintf("sched_setaffinity() is not supported, +setcpuaffinity disabled.\n");
613 }
614
615 #endif