afd5d1a1fc55446c585b8ba2b6b481bd37c061d3
[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)) && !CMK_NO_SOCKETS
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 printf("[%d] %d\n", sysconf(_SC_NPROCESSORS_ONLN), a);
79   return a;
80 }
81
82 /* This implementation assumes the default x86 CPU mask size used by Linux */
83 /* For a large SMP machine, this code should be changed to use a variable sized   */
84 /* CPU affinity mask buffer instead, as the present code will fail beyond 32 CPUs */
85 int set_cpu_affinity(int cpuid) {
86   unsigned long mask = 0xffffffff;
87   unsigned int len = sizeof(mask);
88
89   /* set the affinity mask if possible */
90   if ((cpuid / 8) > len) {
91     printf("Mask size too small to handle requested CPU ID\n");
92     return -1;
93   } else {
94     mask = 1 << cpuid;   /* set the affinity mask exclusively to one CPU */
95   }
96
97 #ifdef _WIN32
98   HANDLE hProcess = GetCurrentProcess();
99   if (SetProcessAffinityMask(hProcess, mask) == 0) {
100     return -1;
101   }
102 #else
103   /* PID 0 refers to the current process */
104   if (sched_setaffinity(0, len, &mask) < 0) {
105     perror("sched_setaffinity");
106     return -1;
107   }
108 #endif
109
110   return 0;
111 }
112
113 int set_thread_affinity(int cpuid) {
114 #if CMK_SMP
115   unsigned long mask = 0xffffffff;
116   unsigned int len = sizeof(mask);
117
118   /* set the affinity mask if possible */
119   if ((cpuid / 8) > len) {
120     printf("Mask size too small to handle requested CPU ID\n");
121     return -1;
122   } else {
123     mask = 1 << cpuid;   /* set the affinity mask exclusively to one CPU */
124   }
125
126 #ifdef _WIN32
127   HANDLE hThread = GetCurrentThread();
128   if (SetThreadAffinityMask(hThread, mask) == 0) {
129     return -1;
130   }
131 #elif  CMK_HAS_PTHREAD_SETAFFINITY
132   /* PID 0 refers to the current process */
133   if (pthread_setaffinity(0, len, &mask) < 0) {
134     perror("pthread_setaffinity");
135     return -1;
136   }
137 #else
138   return set_cpu_affinity(cpuid);
139 #endif
140
141   return 0;
142 #else
143   return -1;
144 #endif
145 }
146
147 /* This implementation assumes the default x86 CPU mask size used by Linux */
148 /* For a large SMP machine, this code should be changed to use a variable sized   */
149 /* CPU affinity mask buffer instead, as the present code will fail beyond 32 CPUs */
150 int print_cpu_affinity() {
151   unsigned long mask;
152   unsigned int len = sizeof(mask);
153
154   /* PID 0 refers to the current process */
155   if (sched_getaffinity(0, len, &mask) < 0) {
156     perror("sched_getaffinity");
157     return -1;
158   }
159
160   printf("CPU affinity mask is: %08lx\n", mask);
161
162   return 0;
163 }
164
165 static int cpuAffinityHandlerIdx;
166 static int cpuAffinityRecvHandlerIdx;
167
168 typedef struct _hostnameMsg {
169   char core[CmiMsgHeaderSizeBytes];
170   int pe;
171   skt_ip_t ip;
172   int ncores;
173   int rank;
174 } hostnameMsg;
175
176 typedef struct _rankMsg {
177   char core[CmiMsgHeaderSizeBytes];
178   int *ranks;
179 } rankMsg;
180
181 static rankMsg *rankmsg = NULL;
182 static CmmTable hostTable;
183
184
185 /* called on PE 0 */
186 static void cpuAffinityHandler(void *m)
187 {
188   static int count = 0;
189   hostnameMsg *rec;
190   hostnameMsg *msg = (hostnameMsg *)m;
191   char str[128];
192   int tag, tag1, pe;
193   CmiAssert(rankmsg != NULL);
194
195   skt_print_ip(str, msg->ip);
196   printf("hostname: %d %s\n", msg->pe, str);
197   tag = *(int*)&msg->ip;
198   pe = msg->pe;
199   if ((rec = (hostnameMsg *)CmmProbe(hostTable, 1, &tag, &tag1)) != NULL) {
200     CmiFree(msg);
201   }
202   else {
203     rec = msg;
204     CmmPut(hostTable, 1, &tag, msg);
205   }
206   rankmsg->ranks[pe] = rec->rank%rec->ncores;
207   rec->rank ++;
208   count ++;
209   if (count == CmiNumPes()) {
210     hostnameMsg *m;
211     tag = CmmWildCard;
212     while (m = CmmGet(hostTable, 1, &tag, &tag1)) CmiFree(m);
213     CmmFree(hostTable);
214     CmiSyncBroadcastAllAndFree(sizeof(rankMsg)+CmiNumPes()*sizeof(int), (void *)rankmsg);
215   }
216 }
217
218 static void cpuAffinityRecvHandler(void *msg)
219 {
220   int myrank;
221   rankMsg *m = (rankMsg *)msg;
222   m->ranks = (int *)((char*)m + sizeof(rankMsg));
223   myrank = m->ranks[CmiMyPe()];
224   CmiPrintf("[%d %d] rank: %d\n", CmiMyNode(), CmiMyPe(), myrank);
225
226   /* set cpu affinity */
227   if (set_cpu_affinity(myrank) != -1)
228     CmiPrintf("Processor %d is bound to core %d\n", CmiMyPe(), myrank);
229   print_cpu_affinity();
230   CmiFree(m);
231 }
232
233 void CmiInitCPUAffinity(char **argv)
234 {
235   skt_ip_t myip;
236   hostnameMsg  *msg;
237   int affinity_flag = CmiGetArgFlagDesc(argv,"+setcpuaffinity",
238                                                 "set cpu affinity");
239
240   cpuAffinityHandlerIdx =
241        CmiRegisterHandler((CmiHandler)cpuAffinityHandler);
242   cpuAffinityRecvHandlerIdx =
243        CmiRegisterHandler((CmiHandler)cpuAffinityRecvHandler);
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   myip = skt_my_ip();
254
255     /* prepare a msg to send */
256   msg = (hostnameMsg *)CmiAlloc(sizeof(hostnameMsg));
257   CmiSetHandler((char *)msg, cpuAffinityHandlerIdx);
258   msg->pe = CmiMyPe();
259   msg->ip = myip;
260   msg->ncores = num_cores();
261   msg->rank = 0;
262   CmiSyncSendAndFree(0, sizeof(hostnameMsg), (void *)msg);
263
264   if (CmiMyPe() == 0) {
265     int i;
266     hostTable = CmmNew();
267     rankmsg = (rankMsg *)CmiAlloc(sizeof(rankMsg)+CmiNumPes()*sizeof(int));
268     CmiSetHandler((char *)rankmsg, cpuAffinityRecvHandlerIdx);
269     rankmsg->ranks = (int *)((char*)rankmsg + sizeof(rankMsg)); /* (int *)(void **)&rankmsg->ranks + 1; */
270     for (i=0; i<CmiNumPes(); i++) rankmsg->ranks[i] = 0;
271   }
272 }
273
274 #else
275
276 void CmiInitCPUAffinity(char **argv)
277 {
278   int affinity_flag = CmiGetArgFlagDesc(argv,"+setcpuaffinity",
279                                                 "set cpu affinity");
280   if (affinity_flag)
281     CmiPrintf("sched_setaffinity() is not supported, +affinity_flag disabled.\n");
282 }
283
284 #endif