lib/nproc.c

   1 /* Detect the number of processors.
   2
   3    Copyright (C) 2009, 2010 Free Software Foundation, Inc.
   4
   5    This program is free software; you can redistribute it and/or modify
   6    it under the terms of the GNU General Public License as published by
   7    the Free Software Foundation; either version 2, or (at your option)
   8    any later version.
   9
  10    This program is distributed in the hope that it will be useful,
  11    but WITHOUT ANY WARRANTY; without even the implied warranty of
  12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13    GNU General Public License for more details.
  14
  15    You should have received a copy of the GNU General Public License
  16    along with this program; if not, write to the Free Software Foundation,
  17    Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.  */
  18
  19 /* Written by Glen Lenker and Bruno Haible.  */
  20
  21 #include <config.h>
  22 #include "nproc.h"
  23
  24 #include <stdlib.h>
  25 #include <unistd.h>
  26
  27 #if HAVE_PTHREAD_AFFINITY_NP && 0
  28 # include <pthread.h>
  29 # include <sched.h>
  30 #endif
  31 #if HAVE_SCHED_GETAFFINITY_LIKE_GLIBC || HAVE_SCHED_GETAFFINITY_NP
  32 # include <sched.h>
  33 #endif
  34
  35 #include <sys/types.h>
  36
  37 #if HAVE_SYS_PSTAT_H
  38 # include <sys/pstat.h>
  39 #endif
  40
  41 #if HAVE_SYS_SYSMP_H
  42 # include <sys/sysmp.h>
  43 #endif
  44
  45 #if HAVE_SYS_PARAM_H
  46 # include <sys/param.h>
  47 #endif
  48
  49 #if HAVE_SYS_SYSCTL_H
  50 # include <sys/sysctl.h>
  51 #endif
  52
  53 #if (defined _WIN32 || defined __WIN32__) && ! defined __CYGWIN__
  54 # define WIN32_LEAN_AND_MEAN
  55 # include <windows.h>
  56 #endif
  57
  58 #include "c-ctype.h"
  59
  60 #define ARRAY_SIZE(a) (sizeof (a) / sizeof ((a)[0]))
  61
  62 unsigned long int
  63 num_processors (enum nproc_query query)
  64 {
  65   if (query == NPROC_CURRENT_OVERRIDABLE)
  66     {
  67       /* Test the environment variable OMP_NUM_THREADS, recognized also by all
  68          programs that are based on OpenMP.  The OpenMP spec says that the
  69          value assigned to the environment variable "may have leading and
  70          trailing white space". */
  71       const char *envvalue = getenv ("OMP_NUM_THREADS");
  72
  73       if (envvalue != NULL)
  74         {
  75           while (*envvalue != '\0' && c_isspace (*envvalue))
  76             envvalue++;
  77           /* Convert it from decimal to 'unsigned long'.  */
  78           if (c_isdigit (*envvalue))
  79             {
  80               char *endptr = NULL;
  81               unsigned long int value = strtoul (envvalue, &endptr, 10);
  82
  83               if (endptr != NULL)
  84                 {
  85                   while (*endptr != '\0' && c_isspace (*endptr))
  86                     endptr++;
  87                   if (*endptr == '\0')
  88                     return (value > 0 ? value : 1);
  89                 }
  90             }
  91         }
  92
  93       query = NPROC_CURRENT;
  94     }
  95   /* Here query is one of NPROC_ALL, NPROC_CURRENT.  */
  96
  97   if (query == NPROC_CURRENT)
  98     {
  99       /* glibc >= 2.3.3 with NPTL and NetBSD 5 have pthread_getaffinity_np,
 100          but with different APIs.  Also it requires linking with -lpthread.
 101          Therefore this code is not enabled.
 102          glibc >= 2.3.4 has sched_getaffinity whereas NetBSD 5 has
 103          sched_getaffinity_np.  */
 104 #if HAVE_PTHREAD_AFFINITY_NP && defined __GLIBC__ && 0
 105       {
 106         cpu_set_t set;
 107
 108         if (pthread_getaffinity_np (pthread_self (), sizeof (set), &set) == 0)
 109           {
 110             unsigned long count;
 111
 112 # ifdef CPU_COUNT
 113             /* glibc >= 2.6 has the CPU_COUNT macro.  */
 114             count = CPU_COUNT (&set);
 115 # else
 116             size_t i;
 117
 118             count = 0;
 119             for (i = 0; i < CPU_SETSIZE; i++)
 120               if (CPU_ISSET (i, &set))
 121                 count++;
 122 # endif
 123             if (count > 0)
 124               return count;
 125           }
 126       }
 127 #elif HAVE_PTHREAD_AFFINITY_NP && defined __NetBSD__ && 0
 128       {
 129         cpuset_t *set;
 130
 131         set = cpuset_create ();
 132         if (set != NULL)
 133           {
 134             unsigned long count = 0;
 135
 136             if (pthread_getaffinity_np (pthread_self (), cpuset_size (set), set)
 137                 == 0)
 138               {
 139                 cpuid_t i;
 140
 141                 for (i = 0;; i++)
 142                   {
 143                     int ret = cpuset_isset (i, set);
 144                     if (ret < 0)
 145                       break;
 146                     if (ret > 0)
 147                       count++;
 148                   }
 149               }
 150             cpuset_destroy (set);
 151             if (count > 0)
 152               return count;
 153           }
 154       }
 155 #elif HAVE_SCHED_GETAFFINITY_LIKE_GLIBC /* glibc >= 2.3.4 */
 156       {
 157         cpu_set_t set;
 158
 159         if (sched_getaffinity (0, sizeof (set), &set) == 0)
 160           {
 161             unsigned long count;
 162
 163 # ifdef CPU_COUNT
 164             /* glibc >= 2.6 has the CPU_COUNT macro.  */
 165             count = CPU_COUNT (&set);
 166 # else
 167             size_t i;
 168
 169             count = 0;
 170             for (i = 0; i < CPU_SETSIZE; i++)
 171               if (CPU_ISSET (i, &set))
 172                 count++;
 173 # endif
 174             if (count > 0)
 175               return count;
 176           }
 177       }
 178 #elif HAVE_SCHED_GETAFFINITY_NP /* NetBSD >= 5 */
 179       {
 180         cpuset_t *set;
 181
 182         set = cpuset_create ();
 183         if (set != NULL)
 184           {
 185             unsigned long count = 0;
 186
 187             if (sched_getaffinity_np (getpid (), cpuset_size (set), set) == 0)
 188               {
 189                 cpuid_t i;
 190
 191                 for (i = 0;; i++)
 192                   {
 193                     int ret = cpuset_isset (i, set);
 194                     if (ret < 0)
 195                       break;
 196                     if (ret > 0)
 197                       count++;
 198                   }
 199               }
 200             cpuset_destroy (set);
 201             if (count > 0)
 202               return count;
 203           }
 204       }
 205 #endif
 206
 207 #if (defined _WIN32 || defined __WIN32__) && ! defined __CYGWIN__
 208       { /* This works on native Windows platforms.  */
 209         DWORD_PTR process_mask;
 210         DWORD_PTR system_mask;
 211
 212         if (GetProcessAffinityMask (GetCurrentProcess (),
 213                                     &process_mask, &system_mask))
 214           {
 215             DWORD_PTR mask = process_mask;
 216             unsigned long count = 0;
 217
 218             for (; mask != 0; mask = mask >> 1)
 219               if (mask & 1)
 220                 count++;
 221             if (count > 0)
 222               return count;
 223           }
 224       }
 225 #endif
 226
 227 #if defined _SC_NPROCESSORS_ONLN
 228       { /* This works on glibc, MacOS X 10.5, FreeBSD, AIX, OSF/1, Solaris,
 229            Cygwin, Haiku.  */
 230         long int nprocs = sysconf (_SC_NPROCESSORS_ONLN);
 231         if (nprocs > 0)
 232           return nprocs;
 233       }
 234 #endif
 235     }
 236   else /* query == NPROC_ALL */
 237     {
 238 #if defined _SC_NPROCESSORS_CONF
 239       { /* This works on glibc, MacOS X 10.5, FreeBSD, AIX, OSF/1, Solaris,
 240            Cygwin, Haiku.  */
 241         long int nprocs = sysconf (_SC_NPROCESSORS_CONF);
 242         if (nprocs > 0)
 243           return nprocs;
 244       }
 245 #endif
 246     }
 247
 248 #if HAVE_PSTAT_GETDYNAMIC
 249   { /* This works on HP-UX.  */
 250     struct pst_dynamic psd;
 251     if (pstat_getdynamic (&psd, sizeof psd, 1, 0) >= 0)
 252       {
 253         /* The field psd_proc_cnt contains the number of active processors.
 254            In newer releases of HP-UX 11, the field psd_max_proc_cnt includes
 255            deactivated processors.  */
 256         if (query == NPROC_CURRENT)
 257           {
 258             if (psd.psd_proc_cnt > 0)
 259               return psd.psd_proc_cnt;
 260           }
 261         else
 262           {
 263             if (psd.psd_max_proc_cnt > 0)
 264               return psd.psd_max_proc_cnt;
 265           }
 266       }
 267   }
 268 #endif
 269
 270 #if HAVE_SYSMP && defined MP_NAPROCS && defined MP_NPROCS
 271   { /* This works on IRIX.  */
 272     /* MP_NPROCS yields the number of installed processors.
 273        MP_NAPROCS yields the number of processors available to unprivileged
 274        processes.  */
 275     int nprocs =
 276       sysmp (query == NPROC_CURRENT && getpid () != 0
 277              ? MP_NAPROCS
 278              : MP_NPROCS);
 279     if (nprocs > 0)
 280       return nprocs;
 281   }
 282 #endif
 283
 284   /* Finally, as fallback, use the APIs that don't distinguish between
 285      NPROC_CURRENT and NPROC_ALL.  */
 286
 287 #if HAVE_SYSCTL && defined HW_NCPU
 288   { /* This works on MacOS X, FreeBSD, NetBSD, OpenBSD.  */
 289     int nprocs;
 290     size_t len = sizeof (nprocs);
 291     static int mib[2] = { CTL_HW, HW_NCPU };
 292
 293     if (sysctl (mib, ARRAY_SIZE (mib), &nprocs, &len, NULL, 0) == 0
 294         && len == sizeof (nprocs)
 295         && 0 < nprocs)
 296       return nprocs;
 297   }
 298 #endif
 299
 300 #if (defined _WIN32 || defined __WIN32__) && ! defined __CYGWIN__
 301   { /* This works on native Windows platforms.  */
 302     SYSTEM_INFO system_info;
 303     GetSystemInfo (&system_info);
 304     if (0 < system_info.dwNumberOfProcessors)
 305       return system_info.dwNumberOfProcessors;
 306   }
 307 #endif
 308
 309   return 1;
 310 }