File tree Expand file tree Collapse file tree 2 files changed +5
-3
lines changed Expand file tree Collapse file tree 2 files changed +5
-3
lines changed Original file line number Diff line number Diff line change @@ -12,6 +12,7 @@ extern "C" {
12
12
/*Set the number of threads on runtime.*/
13
13
void openblas_set_num_threads (int num_threads );
14
14
void goto_set_num_threads (int num_threads );
15
+ int openblas_set_num_threads_local (int num_threads );
15
16
16
17
/*Get the number of threads on runtime.*/
17
18
int openblas_get_num_threads (void );
Original file line number Diff line number Diff line change @@ -137,19 +137,20 @@ typedef struct blas_queue {
137
137
138
138
extern int blas_server_avail ;
139
139
extern int blas_omp_number_max ;
140
+ extern int blas_omp_threads_local ;
140
141
141
142
static __inline int num_cpu_avail (int level ) {
142
143
143
144
#ifdef USE_OPENMP
144
145
int openmp_nthreads ;
145
146
openmp_nthreads = omp_get_max_threads ();
147
+ if (omp_in_parallel ()) openmp_nthreads = blas_omp_threads_local ;
146
148
#endif
147
149
148
150
#ifndef USE_OPENMP
149
151
if (blas_cpu_number == 1
150
- #endif
151
- #ifdef USE_OPENMP
152
- if (openmp_nthreads == 1 || omp_in_parallel ()
152
+ #else
153
+ if (openmp_nthreads == 1
153
154
#endif
154
155
) return 1 ;
155
156
You can’t perform that action at this time.
0 commit comments