use CMK_HAS_GETHOSTNAME to protect code which can not be compiled on bluegene.
[charm.git] / src / conv-core / cpuaffinity.c
1 #include "converse.h"
2 #include "sockRoutines.h"
3
4 /*
5  This scheme relies on using IP address to identify nodes and assigning 
6 cpu affinity.  
7
8  when CMK_NO_SOCKETS, which is typically on cray xt3 and bluegene/L.
9  There is no hostname for the compute nodes.
10 */
11 #if (CMK_HAS_SETAFFINITY || defined (_WIN32)) 
12
13 #include <stdlib.h>
14 #include <stdio.h>
15 #include <unistd.h>
16
17
18 #include <stdlib.h>
19 #include <stdio.h>
20
21 #ifdef _WIN32
22 #include <windows.h>
23 #include <winbase.h>
24 #else
25 #define _GNU_SOURCE
26 #include <sched.h>
27 long sched_setaffinity(pid_t pid, unsigned int len, unsigned long *user_mask_ptr);
28 long sched_getaffinity(pid_t pid, unsigned int len, unsigned long *user_mask_ptr);
29 #endif
30
31 #if defined(__APPLE__) 
32 #include <Carbon/Carbon.h> /* Carbon APIs for Multiprocessing */
33 #endif
34
35 #if defined(ARCH_HPUX11) ||  defined(ARCH_HPUX10)
36 #include <sys/mpctl.h>
37 #endif
38
39
40 int num_cores(void) {
41   int a=1;
42
43   // Allow the user to override the number of CPUs for use
44   // in scalability testing, debugging, etc.
45   char *forcecount = getenv("VMDFORCECPUCOUNT");
46   if (forcecount != NULL) {
47     if (sscanf(forcecount, "%d", &a) == 1) {
48       return a; // if we got a valid count, return it
49     } else {
50       a=1;      // otherwise use the real available hardware CPU count
51     }
52   }
53
54 #if defined(__APPLE__) 
55   a = MPProcessorsScheduled(); /* Number of active/running CPUs */
56 #endif
57
58 #ifdef _WIN32
59   struct _SYSTEM_INFO sysinfo;
60   GetSystemInfo(&sysinfo);
61   a = sysinfo.dwNumberOfProcessors; /* total number of CPUs */
62 #endif /* _MSC_VER */
63
64
65 #ifdef _SC_NPROCESSORS_ONLN
66   a = sysconf(_SC_NPROCESSORS_ONLN); /* number of active/running CPUs */
67 #elif defined(_SC_CRAY_NCPU)
68   a = sysconf(_SC_CRAY_NCPU);
69 #elif defined(_SC_NPROC_ONLN)
70   a = sysconf(_SC_NPROC_ONLN); /* number of active/running CPUs */
71 #endif
72   if (a == -1) a = 1;
73
74 #if defined(ARCH_HPUX11) || defined(ARCH_HPUX10)
75   a = mpctl(MPC_GETNUMSPUS, 0, 0); /* total number of CPUs */
76 #endif /* HPUX */
77
78   return a;
79 }
80
81 /* This implementation assumes the default x86 CPU mask size used by Linux */
82 /* For a large SMP machine, this code should be changed to use a variable sized   */
83 /* CPU affinity mask buffer instead, as the present code will fail beyond 32 CPUs */
84 int set_cpu_affinity(int cpuid) {
85   unsigned long mask = 0xffffffff;
86   unsigned int len = sizeof(mask);
87
88   /* set the affinity mask if possible */
89   if ((cpuid / 8) > len) {
90     printf("Mask size too small to handle requested CPU ID\n");
91     return -1;
92   } else {
93     mask = 1 << cpuid;   /* set the affinity mask exclusively to one CPU */
94   }
95
96 #ifdef _WIN32
97   HANDLE hProcess = GetCurrentProcess();
98   if (SetProcessAffinityMask(hProcess, mask) == 0) {
99     return -1;
100   }
101 #else
102   /* PID 0 refers to the current process */
103   if (sched_setaffinity(0, len, &mask) < 0) {
104     perror("sched_setaffinity");
105     return -1;
106   }
107 #endif
108
109   return 0;
110 }
111
112 int set_thread_affinity(int cpuid) {
113 #if CMK_SMP
114   unsigned long mask = 0xffffffff;
115   unsigned int len = sizeof(mask);
116
117   /* set the affinity mask if possible */
118   if ((cpuid / 8) > len) {
119     printf("Mask size too small to handle requested CPU ID\n");
120     return -1;
121   } else {
122     mask = 1 << cpuid;   /* set the affinity mask exclusively to one CPU */
123   }
124
125 #ifdef _WIN32
126   HANDLE hThread = GetCurrentThread();
127   if (SetThreadAffinityMask(hThread, mask) == 0) {
128     return -1;
129   }
130 #elif  CMK_HAS_PTHREAD_SETAFFINITY
131   /* PID 0 refers to the current process */
132   if (pthread_setaffinity(0, len, &mask) < 0) {
133     perror("pthread_setaffinity");
134     return -1;
135   }
136 #else
137   return set_cpu_affinity(cpuid);
138 #endif
139
140   return 0;
141 #else
142   return -1;
143 #endif
144 }
145
146 /* This implementation assumes the default x86 CPU mask size used by Linux */
147 /* For a large SMP machine, this code should be changed to use a variable sized   */
148 /* CPU affinity mask buffer instead, as the present code will fail beyond 32 CPUs */
149 int print_cpu_affinity() {
150   unsigned long mask;
151   unsigned int len = sizeof(mask);
152
153   /* PID 0 refers to the current process */
154   if (sched_getaffinity(0, len, &mask) < 0) {
155     perror("sched_getaffinity");
156     return -1;
157   }
158
159   printf("CPU affinity mask is: %08lx\n", mask);
160
161   return 0;
162 }
163
164 static int cpuAffinityHandlerIdx;
165 static int cpuAffinityRecvHandlerIdx;
166
167 typedef struct _hostnameMsg {
168   char core[CmiMsgHeaderSizeBytes];
169   int pe;
170   skt_ip_t ip;
171   int ncores;
172   int rank;
173 } hostnameMsg;
174
175 typedef struct _rankMsg {
176   char core[CmiMsgHeaderSizeBytes];
177   int *ranks;
178 } rankMsg;
179
180 static rankMsg *rankmsg = NULL;
181 static CmmTable hostTable;
182
183
184 /* called on PE 0 */
185 static void cpuAffinityHandler(void *m)
186 {
187   static int count = 0;
188   hostnameMsg *rec;
189   hostnameMsg *msg = (hostnameMsg *)m;
190   char str[128];
191   int tag, tag1, pe;
192   CmiAssert(rankmsg != NULL);
193
194   skt_print_ip(str, msg->ip);
195   printf("hostname: %d %s\n", msg->pe, str);
196   tag = *(int*)&msg->ip;
197   pe = msg->pe;
198   if ((rec = (hostnameMsg *)CmmProbe(hostTable, 1, &tag, &tag1)) != NULL) {
199     CmiFree(msg);
200   }
201   else {
202     rec = msg;
203     CmmPut(hostTable, 1, &tag, msg);
204   }
205   rankmsg->ranks[pe] = rec->rank%rec->ncores;
206   rec->rank ++;
207   count ++;
208   if (count == CmiNumPes()) {
209     hostnameMsg *m;
210     tag = CmmWildCard;
211     while (m = CmmGet(hostTable, 1, &tag, &tag1)) CmiFree(m);
212     CmmFree(hostTable);
213     CmiSyncBroadcastAllAndFree(sizeof(rankMsg)+CmiNumPes()*sizeof(int), (void *)rankmsg);
214   }
215 }
216
217 static void cpuAffinityRecvHandler(void *msg)
218 {
219   int myrank;
220   rankMsg *m = (rankMsg *)msg;
221   m->ranks = (int *)((char*)m + sizeof(rankMsg));
222   myrank = m->ranks[CmiMyPe()];
223   CmiPrintf("[%d %d] rank: %d\n", CmiMyNode(), CmiMyPe(), myrank);
224
225   /* set cpu affinity */
226   if (set_cpu_affinity(myrank) != -1)
227     CmiPrintf("Processor %d is bound to core %d\n", CmiMyPe(), myrank);
228   print_cpu_affinity();
229   CmiFree(m);
230 }
231
232 void CmiInitCPUAffinity(char **argv)
233 {
234   skt_ip_t myip;
235   hostnameMsg  *msg;
236   int affinity_flag = CmiGetArgFlagDesc(argv,"+setcpuaffinity",
237                                                 "set cpu affinity");
238
239   cpuAffinityHandlerIdx =
240        CmiRegisterHandler((CmiHandler)cpuAffinityHandler);
241   cpuAffinityRecvHandlerIdx =
242        CmiRegisterHandler((CmiHandler)cpuAffinityRecvHandler);
243
244   if (CmiMyPe() >= CmiNumPes()) return;    /* comm thread return */
245   if (!affinity_flag) return;
246
247 #if 0
248   if (gethostname(hostname, 999)!=0) {
249       strcpy(hostname, "");
250   }
251 #endif
252     /* get my ip address */
253 #if CMK_XT3
254   myip = getXT3NodeID(CmiMyPe(), CmiNumPes());
255 #elif CMK_HAS_GETHOSTNAME
256   myip = skt_my_ip();
257 #else
258   CmiAbort("Can not get unique name for the compute nodes. \n");
259 #endif
260
261     /* prepare a msg to send */
262   msg = (hostnameMsg *)CmiAlloc(sizeof(hostnameMsg));
263   CmiSetHandler((char *)msg, cpuAffinityHandlerIdx);
264   msg->pe = CmiMyPe();
265   msg->ip = myip;
266   msg->ncores = num_cores();
267   msg->rank = 0;
268   CmiSyncSendAndFree(0, sizeof(hostnameMsg), (void *)msg);
269
270   if (CmiMyPe() == 0) {
271     int i;
272     hostTable = CmmNew();
273     rankmsg = (rankMsg *)CmiAlloc(sizeof(rankMsg)+CmiNumPes()*sizeof(int));
274     CmiSetHandler((char *)rankmsg, cpuAffinityRecvHandlerIdx);
275     rankmsg->ranks = (int *)((char*)rankmsg + sizeof(rankMsg)); /* (int *)(void **)&rankmsg->ranks + 1; */
276     for (i=0; i<CmiNumPes(); i++) rankmsg->ranks[i] = 0;
277   }
278 }
279
280 #else
281
282 void CmiInitCPUAffinity(char **argv)
283 {
284   int affinity_flag = CmiGetArgFlagDesc(argv,"+setcpuaffinity",
285                                                 "set cpu affinity");
286   if (affinity_flag)
287     CmiPrintf("sched_setaffinity() is not supported, +affinity_flag disabled.\n");
288 }
289
290 #endif