From 41137bbc339ede60f39eaf9db3d61188e0ab1e8c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A1draig=20Brady?= <address@hidden>
Date: Mon, 20 Feb 2017 20:26:35 -0800
Subject: [PATCH] nproc: support nested OMP_NUM_THREADS, and OMP_THREAD_LIMIT

* lib/nproc.c (parse_omp_threads): A new function refactored
from num_processors() to support parsing both of the
above environment variables.
(num_processors): Prefer using omp_get_num_threads() with [_OPENMP]
to accurately reflect the current OpenMP nesting level.
Also support the OMP_THREAD_LIMIT environment variable
to limit the max value determined from OMP_NUM_THREADS.
* modules/nproc: Depend on minmax header.
Suggested by Oliver Heimlich.
---
 lib/nproc.c   | 82 +++++++++++++++++++++++++++++++++++++++++++----------------
 modules/nproc |  1 +
 2 files changed, 61 insertions(+), 22 deletions(-)

diff --git a/lib/nproc.c b/lib/nproc.c
index 9053d75..3d9007b 100644
--- a/lib/nproc.c
+++ b/lib/nproc.c
@@ -54,8 +54,14 @@
 # include <windows.h>
 #endif
 
+#if defined(_OPENMP)
+# include <omp.h>
+#endif
+
 #include "c-ctype.h"
 
+#include "minmax.h"
+
 #define ARRAY_SIZE(a) (sizeof (a) / sizeof ((a)[0]))
 
 /* Return the number of processors available to the current process, based
@@ -196,35 +202,67 @@ num_processors_via_affinity_mask (void)
   return 0;
 }
 
+/* Parse OMP environment variables without dependence on OMP.
+   Return > 0 for valid values.  */
+static unsigned long int
+parse_omp_threads (char const* threads)
+{
+  unsigned long int ret = 0;
+
+  if (threads == NULL)
+    return ret;
+
+  /* The OpenMP spec says that the value assigned to the environment variables
+     "may have leading and trailing white space".  */
+  while (*threads != '\0' && c_isspace (*threads))
+    threads++;
+
+  /* Convert it from decimal to 'unsigned long'.  */
+  if (c_isdigit (*threads))
+    {
+      char *endptr = NULL;
+      unsigned long int value = strtoul (threads, &endptr, 10);
+
+      if (endptr != NULL)
+        {
+          while (*endptr != '\0' && c_isspace (*endptr))
+            endptr++;
+          if (*endptr == '\0')
+            return (value > 0 ? value : 1);
+          /* Also accept the first value in a nesting level,
+             since we can't determine the nesting level from env vars.  */
+          else if (*endptr == ',')
+            return (value > 0 ? value : 1);
+        }
+    }
+
+  return ret;
+}
+
 unsigned long int
 num_processors (enum nproc_query query)
 {
   if (query == NPROC_CURRENT_OVERRIDABLE)
     {
-      /* Test the environment variable OMP_NUM_THREADS, recognized also by all
-         programs that are based on OpenMP.  The OpenMP spec says that the
-         value assigned to the environment variable "may have leading and
-         trailing white space". */
-      const char *envvalue = getenv ("OMP_NUM_THREADS");
+      unsigned long int omp_env_threads;
+      unsigned long int omp_env_limit;
+
+/* If OPENMP is enabled with AC_OPENMP, then use OPENMP to
+   get the number of threads for the current nesting level.  */
+#if defined(_OPENMP)
+      return omp_get_num_threads ();
+#endif
 
-      if (envvalue != NULL)
+      /* Honor the OpenMP environment variables, recognized also by all
+         programs that are based on OpenMP.  */
+      omp_env_threads = parse_omp_threads (getenv ("OMP_NUM_THREADS"));
+      if (omp_env_threads)
         {
-          while (*envvalue != '\0' && c_isspace (*envvalue))
-            envvalue++;
-          /* Convert it from decimal to 'unsigned long'.  */
-          if (c_isdigit (*envvalue))
-            {
-              char *endptr = NULL;
-              unsigned long int value = strtoul (envvalue, &endptr, 10);
-
-              if (endptr != NULL)
-                {
-                  while (*endptr != '\0' && c_isspace (*endptr))
-                    endptr++;
-                  if (*endptr == '\0')
-                    return (value > 0 ? value : 1);
-                }
-            }
+          omp_env_limit = parse_omp_threads (getenv ("OMP_THREADS_LIMIT"));
+          if (omp_env_limit > 1)
+            omp_env_threads = MIN (omp_env_threads, omp_env_limit);
+
+          return omp_env_threads;
         }
 
       query = NPROC_CURRENT;
diff --git a/modules/nproc b/modules/nproc
index 2cbf04d..5259579 100644
--- a/modules/nproc
+++ b/modules/nproc
@@ -9,6 +9,7 @@ m4/nproc.m4
 Depends-on:
 c-ctype
 extensions
+minmax
 unistd
 
 configure.ac:
-- 
2.5.5