Use sysconf to detect cpu count

1 files changed, 20 insertions, 13 deletions

diff --git a/threadpool.c b/threadpool.c
index fb6f0ee..fb93a18 100644
--- a/threadpool.c
+++ b/threadpool.c
@@ -19,6 +19,9 @@
 #include <omp.h>
 #include <stdlib.h>
 #include <string.h>
+#include <unistd.h>
+
+#include <stdio.h>
 
 #include "threadpool.h"
 
@@ -40,12 +43,25 @@ void tp_free(TPContext **pctx)
 int tp_init(TPContext **pctx, unsigned int nb_threads)
 {
     TPContext *ctx = NULL;
-    char *env_threads;
     int ret;
 
     if (!nb_threads) {
-        ret = -EINVAL;
-        goto fail;
+        const char *env_threads = getenv("OMP_NUM_THREADS");
+        if (env_threads) {
+            nb_threads = strtol(env_threads, NULL, 0);
+        }
+#ifdef _SC_NPROCESSORS_ONLN
+        else {
+            long val = sysconf(_SC_NPROCESSORS_ONLN);
+            if (val > 0)
+                nb_threads = val;
+        }
+#endif
+
+        if (!nb_threads) {
+            ret = -EINVAL;
+            goto fail;
+        }
     }
 
     ctx = calloc(1, sizeof(*ctx));
@@ -54,16 +70,7 @@ int tp_init(TPContext **pctx, unsigned int nb_threads)
         goto fail;
     }
 
-    env_threads = getenv("OMP_NUM_THREADS");
-    if (env_threads) {
-        nb_threads = strtol(env_threads, NULL, 0);
-        if (!nb_threads) {
-            ret = -EINVAL;
-            goto fail;
-        }
-    } else {
-        omp_set_num_threads(nb_threads);
-    }
+    omp_set_num_threads(nb_threads);
 
     ctx->nb_threads = nb_threads;