Skip to content

Commit f8674a7

Browse files
committed
Make OpenBLAS's usage of OpenMP respect openblas_set_num_threads(). Fixes #5806.
Until now, the code in `num_cpu_avail()`, if (blas_cpu_number != openmp_nthreads) { goto_set_num_threads(openmp_nthreads); } would just always set the threads back to OpenMP's thread count.
1 parent 8cecf89 commit f8674a7

2 files changed

Lines changed: 14 additions & 4 deletions

File tree

common_thread.h

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,10 +138,19 @@ typedef struct blas_queue {
138138
extern int blas_server_avail;
139139
extern int blas_omp_number_max;
140140
extern int blas_omp_threads_local;
141+
extern int blas_is_num_threads_set_explicitly;
141142

142143
static __inline int num_cpu_avail(int level) {
143144

144145
#ifdef USE_OPENMP
146+
/* If the user explicitly called openblas_set_num_threads(),
147+
respect that setting instead of overriding it with
148+
`omp_get_max_threads()` below (which is to get a default
149+
in case the user hasn't made an explicit choice). */
150+
if (blas_is_num_threads_set_explicitly) {
151+
return blas_cpu_number;
152+
}
153+
145154
int openmp_nthreads;
146155
openmp_nthreads=omp_get_max_threads();
147156
if (omp_in_parallel()) openmp_nthreads = blas_omp_threads_local;

driver/others/blas_server_omp.c

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
int blas_server_avail = 0;
7171
int blas_omp_number_max = 0;
7272
int blas_omp_threads_local = 1;
73+
int blas_is_num_threads_set_explicitly = 0; // tracks whether the user called openblas_set_num_threads()
7374

7475
extern int openblas_omp_adaptive_env(void);
7576

@@ -122,7 +123,7 @@ void goto_set_num_threads(int num_threads) {
122123

123124
}
124125
void openblas_set_num_threads(int num_threads) {
125-
126+
blas_is_num_threads_set_explicitly = 1;
126127
goto_set_num_threads(num_threads);
127128
}
128129

@@ -145,7 +146,7 @@ extern int openblas_omp_num_threads_env(void);
145146

146147
if(blas_omp_number_max <= 0)
147148
blas_omp_number_max= openblas_omp_num_threads_env();
148-
if (blas_omp_number_max <= 0)
149+
if (blas_omp_number_max <= 0)
149150
blas_omp_number_max=MAX_CPU_NUMBER;
150151
#else
151152
blas_omp_number_max = omp_get_max_threads();
@@ -365,14 +366,14 @@ static void exec_threads(int thread_num, blas_queue_t *queue, int buf_index){
365366
#ifdef BUILD_COMPLEX16
366367
sb = (void *)(((BLASLONG)sa + ((ZGEMM_P * ZGEMM_Q * 2 * sizeof(double)
367368
+ GEMM_ALIGN) & ~GEMM_ALIGN)) + GEMM_OFFSET_B);
368-
#else
369+
#else
369370
fprintf(stderr,"UNHANDLED COMPLEX16\n");
370371
#endif
371372
} else if ((queue -> mode & BLAS_PREC) == BLAS_SINGLE) {
372373
#ifdef BUILD_COMPLEX
373374
sb = (void *)(((BLASLONG)sa + ((CGEMM_P * CGEMM_Q * 2 * sizeof(float)
374375
+ GEMM_ALIGN) & ~GEMM_ALIGN)) + GEMM_OFFSET_B);
375-
#else
376+
#else
376377
fprintf(stderr,"UNHANDLED COMPLEX\n");
377378
#endif
378379
} else {

0 commit comments

Comments
 (0)