Merge pull request #5808 from nh2/issue-5806-openblas_set_num_threads-openmp

Make OpenBLAS's usage of OpenMP respect `openblas_set_num_threads()`

Author: martin-frbg

common_thread.h

@@ -138,10 +138,19 @@ typedef struct blas_queue {
 extern int blas_server_avail;
 extern int blas_omp_number_max;
 extern int blas_omp_threads_local;
+extern int blas_is_num_threads_set_explicitly;
 
 static __inline int num_cpu_avail(int level) {
 
 #ifdef USE_OPENMP
+  /* If the user explicitly called openblas_set_num_threads(),
+     respect that setting instead of overriding it with
+     `omp_get_max_threads()` below (which is to get a default
+     in case the user hasn't made an explicit choice). */
+  if (blas_is_num_threads_set_explicitly) {
+    return blas_cpu_number;
+  }
+
 int openmp_nthreads;
 	openmp_nthreads=omp_get_max_threads();
 	if (omp_in_parallel()) openmp_nthreads = blas_omp_threads_local;

driver/others/blas_server_omp.c

@@ -70,6 +70,7 @@
 int blas_server_avail = 0;
 int blas_omp_number_max = 0;
 int blas_omp_threads_local = 1;
+int blas_is_num_threads_set_explicitly = 0; // tracks whether the user called openblas_set_num_threads()
 
 extern int openblas_omp_adaptive_env(void);
 
@@ -118,7 +119,7 @@ void goto_set_num_threads(int num_threads) {
 
 }
 void openblas_set_num_threads(int num_threads) {
-
+	blas_is_num_threads_set_explicitly = 1;
 	goto_set_num_threads(num_threads);
 }
 
@@ -141,7 +142,7 @@ extern int openblas_omp_num_threads_env(void);
 
    if(blas_omp_number_max <= 0)
 	   blas_omp_number_max= openblas_omp_num_threads_env();
-   if (blas_omp_number_max <= 0) 
+   if (blas_omp_number_max <= 0)
 	   blas_omp_number_max=MAX_CPU_NUMBER;
 #else
     blas_omp_number_max = omp_get_max_threads();
@@ -361,14 +362,14 @@ static void exec_threads(int thread_num, blas_queue_t *queue, int buf_index){
 #ifdef BUILD_COMPLEX16
 	    sb = (void *)(((BLASLONG)sa + ((ZGEMM_P * ZGEMM_Q * 2 * sizeof(double)
 					    + GEMM_ALIGN) & ~GEMM_ALIGN)) + GEMM_OFFSET_B);
-#else 
+#else
 fprintf(stderr,"UNHANDLED COMPLEX16\n");
 #endif
 	  } else if ((queue -> mode & BLAS_PREC) == BLAS_SINGLE) {
 #ifdef BUILD_COMPLEX
 	    sb = (void *)(((BLASLONG)sa + ((CGEMM_P * CGEMM_Q * 2 * sizeof(float)
 					    + GEMM_ALIGN) & ~GEMM_ALIGN)) + GEMM_OFFSET_B);
-#else 
+#else
 fprintf(stderr,"UNHANDLED COMPLEX\n");
 #endif
 	  } else {