lib/nproc.c

   1 /* Detect the number of processors.
   2
   3    Copyright (C) 2009-2019 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, see <https://www.gnu.org/licenses/>.  */
  17
  18 /* Written by Glen Lenker and Bruno Haible.  */
  19
  20 #include <config.h>
  21 #include "nproc.h"
  22
  23 #include <limits.h>
  24 #include <stdlib.h>
  25 #include <unistd.h>
  26
  27 #if HAVE_PTHREAD_GETAFFINITY_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 __CYGWIN__
  54 # define WIN32_LEAN_AND_MEAN
  55 # include <windows.h>
  56 #endif
  57
  58 #include "c-ctype.h"
  59
  60 #include "minmax.h"
  61
  62 #define ARRAY_SIZE(a) (sizeof (a) / sizeof ((a)[0]))
  63
  64 /* Return the number of processors available to the current process, based
  65    on a modern system call that returns the "affinity" between the current
  66    process and each CPU.  Return 0 if unknown or if such a system call does
  67    not exist.  */
  68 static unsigned long
  69 num_processors_via_affinity_mask (void)
  70 {
  71   /* glibc >= 2.3.3 with NPTL and NetBSD 5 have pthread_getaffinity_np,
  72      but with different APIs.  Also it requires linking with -lpthread.
  73      Therefore this code is not enabled.
  74      glibc >= 2.3.4 has sched_getaffinity whereas NetBSD 5 has
  75      sched_getaffinity_np.  */
  76 #if HAVE_PTHREAD_GETAFFINITY_NP && defined __GLIBC__ && 0
  77   {
  78     cpu_set_t set;
  79
  80     if (pthread_getaffinity_np (pthread_self (), sizeof (set), &set) == 0)
  81       {
  82         unsigned long count;
  83
  84 # ifdef CPU_COUNT
  85         /* glibc >= 2.6 has the CPU_COUNT macro.  */
  86         count = CPU_COUNT (&set);
  87 # else
  88         size_t i;
  89
  90         count = 0;
  91         for (i = 0; i < CPU_SETSIZE; i++)
  92           if (CPU_ISSET (i, &set))
  93             count++;
  94 # endif
  95         if (count > 0)
  96           return count;
  97       }
  98   }
  99 #elif HAVE_PTHREAD_GETAFFINITY_NP && defined __NetBSD__ && 0
 100   {
 101     cpuset_t *set;
 102
 103     set = cpuset_create ();
 104     if (set != NULL)
 105       {
 106         unsigned long count = 0;
 107
 108         if (pthread_getaffinity_np (pthread_self (), cpuset_size (set), set)
 109             == 0)
 110           {
 111             cpuid_t i;
 112
 113             for (i = 0;; i++)
 114               {
 115                 int ret = cpuset_isset (i, set);
 116                 if (ret < 0)
 117                   break;
 118                 if (ret > 0)
 119                   count++;
 120               }
 121           }
 122         cpuset_destroy (set);
 123         if (count > 0)
 124           return count;
 125       }
 126   }
 127 #elif HAVE_SCHED_GETAFFINITY_LIKE_GLIBC /* glibc >= 2.3.4 */
 128   {
 129     cpu_set_t set;
 130
 131     if (sched_getaffinity (0, sizeof (set), &set) == 0)
 132       {
 133         unsigned long count;
 134
 135 # ifdef CPU_COUNT
 136         /* glibc >= 2.6 has the CPU_COUNT macro.  */
 137         count = CPU_COUNT (&set);
 138 # else
 139         size_t i;
 140
 141         count = 0;
 142         for (i = 0; i < CPU_SETSIZE; i++)
 143           if (CPU_ISSET (i, &set))
 144             count++;
 145 # endif
 146         if (count > 0)
 147           return count;
 148       }
 149   }
 150 #elif HAVE_SCHED_GETAFFINITY_NP /* NetBSD >= 5 */
 151   {
 152     cpuset_t *set;
 153
 154     set = cpuset_create ();
 155     if (set != NULL)
 156       {
 157         unsigned long count = 0;
 158
 159         if (sched_getaffinity_np (getpid (), cpuset_size (set), set) == 0)
 160           {
 161             cpuid_t i;
 162
 163             for (i = 0;; i++)
 164               {
 165                 int ret = cpuset_isset (i, set);
 166                 if (ret < 0)
 167                   break;
 168                 if (ret > 0)
 169                   count++;
 170               }
 171           }
 172         cpuset_destroy (set);
 173         if (count > 0)
 174           return count;
 175       }
 176   }
 177 #endif
 178
 179 #if defined _WIN32 && ! defined __CYGWIN__
 180   { /* This works on native Windows platforms.  */
 181     DWORD_PTR process_mask;
 182     DWORD_PTR system_mask;
 183
 184     if (GetProcessAffinityMask (GetCurrentProcess (),
 185                                 &process_mask, &system_mask))
 186       {
 187         DWORD_PTR mask = process_mask;
 188         unsigned long count = 0;
 189
 190         for (; mask != 0; mask = mask >> 1)
 191           if (mask & 1)
 192             count++;
 193         if (count > 0)
 194           return count;
 195       }
 196   }
 197 #endif
 198
 199   return 0;
 200 }
 201
 202
 203 /* Return the total number of processors.  Here QUERY must be one of
 204    NPROC_ALL, NPROC_CURRENT.  The result is guaranteed to be at least 1.  */
 205 static unsigned long int
 206 num_processors_ignoring_omp (enum nproc_query query)
 207 {
 208   /* On systems with a modern affinity mask system call, we have
 209          sysconf (_SC_NPROCESSORS_CONF)
 210             >= sysconf (_SC_NPROCESSORS_ONLN)
 211                >= num_processors_via_affinity_mask ()
 212      The first number is the number of CPUs configured in the system.
 213      The second number is the number of CPUs available to the scheduler.
 214      The third number is the number of CPUs available to the current process.
 215
 216      Note! On Linux systems with glibc, the first and second number come from
 217      the /sys and /proc file systems (see
 218      glibc/sysdeps/unix/sysv/linux/getsysstats.c).
 219      In some situations these file systems are not mounted, and the sysconf
 220      call returns 1, which does not reflect the reality.  */
 221
 222   if (query == NPROC_CURRENT)
 223     {
 224       /* Try the modern affinity mask system call.  */
 225       {
 226         unsigned long nprocs = num_processors_via_affinity_mask ();
 227
 228         if (nprocs > 0)
 229           return nprocs;
 230       }
 231
 232 #if defined _SC_NPROCESSORS_ONLN
 233       { /* This works on glibc, Mac OS X 10.5, FreeBSD, AIX, OSF/1, Solaris,
 234            Cygwin, Haiku.  */
 235         long int nprocs = sysconf (_SC_NPROCESSORS_ONLN);
 236         if (nprocs > 0)
 237           return nprocs;
 238       }
 239 #endif
 240     }
 241   else /* query == NPROC_ALL */
 242     {
 243 #if defined _SC_NPROCESSORS_CONF
 244       { /* This works on glibc, Mac OS X 10.5, FreeBSD, AIX, OSF/1, Solaris,
 245            Cygwin, Haiku.  */
 246         long int nprocs = sysconf (_SC_NPROCESSORS_CONF);
 247
 248 # if __GLIBC__ >= 2 && defined __linux__
 249         /* On Linux systems with glibc, this information comes from the /sys and
 250            /proc file systems (see glibc/sysdeps/unix/sysv/linux/getsysstats.c).
 251            In some situations these file systems are not mounted, and the
 252            sysconf call returns 1.  But we wish to guarantee that
 253            num_processors (NPROC_ALL) >= num_processors (NPROC_CURRENT).  */
 254         if (nprocs == 1)
 255           {
 256             unsigned long nprocs_current = num_processors_via_affinity_mask ();
 257
 258             if (nprocs_current > 0)
 259               nprocs = nprocs_current;
 260           }
 261 # endif
 262
 263         if (nprocs > 0)
 264           return nprocs;
 265       }
 266 #endif
 267     }
 268
 269 #if HAVE_PSTAT_GETDYNAMIC
 270   { /* This works on HP-UX.  */
 271     struct pst_dynamic psd;
 272     if (pstat_getdynamic (&psd, sizeof psd, 1, 0) >= 0)
 273       {
 274         /* The field psd_proc_cnt contains the number of active processors.
 275            In newer releases of HP-UX 11, the field psd_max_proc_cnt includes
 276            deactivated processors.  */
 277         if (query == NPROC_CURRENT)
 278           {
 279             if (psd.psd_proc_cnt > 0)
 280               return psd.psd_proc_cnt;
 281           }
 282         else
 283           {
 284             if (psd.psd_max_proc_cnt > 0)
 285               return psd.psd_max_proc_cnt;
 286           }
 287       }
 288   }
 289 #endif
 290
 291 #if HAVE_SYSMP && defined MP_NAPROCS && defined MP_NPROCS
 292   { /* This works on IRIX.  */
 293     /* MP_NPROCS yields the number of installed processors.
 294        MP_NAPROCS yields the number of processors available to unprivileged
 295        processes.  */
 296     int nprocs =
 297       sysmp (query == NPROC_CURRENT && getpid () != 0
 298              ? MP_NAPROCS
 299              : MP_NPROCS);
 300     if (nprocs > 0)
 301       return nprocs;
 302   }
 303 #endif
 304
 305   /* Finally, as fallback, use the APIs that don't distinguish between
 306      NPROC_CURRENT and NPROC_ALL.  */
 307
 308 #if HAVE_SYSCTL && defined HW_NCPU
 309   { /* This works on Mac OS X, FreeBSD, NetBSD, OpenBSD.  */
 310     int nprocs;
 311     size_t len = sizeof (nprocs);
 312     static int mib[2] = { CTL_HW, HW_NCPU };
 313
 314     if (sysctl (mib, ARRAY_SIZE (mib), &nprocs, &len, NULL, 0) == 0
 315         && len == sizeof (nprocs)
 316         && 0 < nprocs)
 317       return nprocs;
 318   }
 319 #endif
 320
 321 #if defined _WIN32 && ! defined __CYGWIN__
 322   { /* This works on native Windows platforms.  */
 323     SYSTEM_INFO system_info;
 324     GetSystemInfo (&system_info);
 325     if (0 < system_info.dwNumberOfProcessors)
 326       return system_info.dwNumberOfProcessors;
 327   }
 328 #endif
 329
 330   return 1;
 331 }
 332
 333 /* Parse OMP environment variables without dependence on OMP.
 334    Return 0 for invalid values.  */
 335 static unsigned long int
 336 parse_omp_threads (char const* threads)
 337 {
 338   unsigned long int ret = 0;
 339
 340   if (threads == NULL)
 341     return ret;
 342
 343   /* The OpenMP spec says that the value assigned to the environment variables
 344      "may have leading and trailing white space".  */
 345   while (*threads != '\0' && c_isspace (*threads))
 346     threads++;
 347
 348   /* Convert it from positive decimal to 'unsigned long'.  */
 349   if (c_isdigit (*threads))
 350     {
 351       char *endptr = NULL;
 352       unsigned long int value = strtoul (threads, &endptr, 10);
 353
 354       if (endptr != NULL)
 355         {
 356           while (*endptr != '\0' && c_isspace (*endptr))
 357             endptr++;
 358           if (*endptr == '\0')
 359             return value;
 360           /* Also accept the first value in a nesting level,
 361              since we can't determine the nesting level from env vars.  */
 362           else if (*endptr == ',')
 363             return value;
 364         }
 365     }
 366
 367   return ret;
 368 }
 369
 370 unsigned long int
 371 num_processors (enum nproc_query query)
 372 {
 373   unsigned long int omp_env_limit = ULONG_MAX;
 374
 375   if (query == NPROC_CURRENT_OVERRIDABLE)
 376     {
 377       unsigned long int omp_env_threads;
 378       /* Honor the OpenMP environment variables, recognized also by all
 379          programs that are based on OpenMP.  */
 380       omp_env_threads = parse_omp_threads (getenv ("OMP_NUM_THREADS"));
 381       omp_env_limit = parse_omp_threads (getenv ("OMP_THREAD_LIMIT"));
 382       if (! omp_env_limit)
 383         omp_env_limit = ULONG_MAX;
 384
 385       if (omp_env_threads)
 386         return MIN (omp_env_threads, omp_env_limit);
 387
 388       query = NPROC_CURRENT;
 389     }
 390   /* Here query is one of NPROC_ALL, NPROC_CURRENT.  */
 391   {
 392     unsigned long nprocs = num_processors_ignoring_omp (query);
 393     return MIN (nprocs, omp_env_limit);
 394   }
 395 }