Merge pull request #4425 from martin-frbg/issue2392

Add BLAS extension openblas_set_num_threads_local()

Author: martin-frbg

cblas.h

@@ -12,6 +12,7 @@ extern "C" {
 /*Set the number of threads on runtime.*/
 void openblas_set_num_threads(int num_threads);
 void goto_set_num_threads(int num_threads);
+int openblas_set_num_threads_local(int num_threads);
 
 /*Get the number of threads on runtime.*/
 int openblas_get_num_threads(void);

common_thread.h

@@ -137,19 +137,20 @@ typedef struct blas_queue {
 
 extern int blas_server_avail;
 extern int blas_omp_number_max;
+extern int blas_omp_threads_local;
 
 static __inline int num_cpu_avail(int level) {
 
 #ifdef USE_OPENMP
 int openmp_nthreads;
 	openmp_nthreads=omp_get_max_threads();
+	if (omp_in_parallel()) openmp_nthreads = blas_omp_threads_local;
 #endif
 
 #ifndef USE_OPENMP 
   if (blas_cpu_number == 1
-#endif
-#ifdef USE_OPENMP
-     if (openmp_nthreads == 1 || omp_in_parallel()
+#else
+     if (openmp_nthreads == 1 
 #endif
       ) return 1;        
 

driver/others/blas_server.c

@@ -113,6 +113,8 @@ extern unsigned int openblas_thread_timeout(void);
 /* We need this global for checking if initialization is finished.  */
 int blas_server_avail   __attribute__((aligned(ATTRIBUTE_SIZE))) = 0;
 
+int blas_omp_threads_local = 1;
+
 /* Local Variables */
 #if   defined(USE_PTHREAD_LOCK)
 static pthread_mutex_t  server_lock    = PTHREAD_MUTEX_INITIALIZER;

driver/others/blas_server_omp.c

@@ -69,6 +69,7 @@
 
 int blas_server_avail = 0;
 int blas_omp_number_max = 0;
+int blas_omp_threads_local = 1;
 
 extern int openblas_omp_adaptive_env(void);
 

driver/others/blas_server_win32.c

@@ -59,6 +59,8 @@ static CRITICAL_SECTION queue_lock;
 /* We need this global for checking if initialization is finished.   */
 int blas_server_avail = 0;
 
+int blas_omp_threads_local = 1;
+
 /* Local Variables */
 static BLASULONG server_lock       = 0;
 

driver/others/openblas_set_num_threads.c

@@ -36,11 +36,20 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #ifdef SMP_SERVER
 
 extern  void openblas_set_num_threads(int num_threads) ;
+extern  int openblas_get_num_threads(void) ;
 
 void openblas_set_num_threads_(int* num_threads){
 	openblas_set_num_threads(*num_threads);
 }
 
+int openblas_set_num_threads_local(int num_threads){
+	int ret = openblas_get_num_threads();
+	openblas_set_num_threads(num_threads);
+	blas_omp_threads_local=num_threads;
+	return ret;
+}
+
+
 #else
 //Single thread
 
@@ -50,4 +59,8 @@ void openblas_set_num_threads(int num_threads) {
 void openblas_set_num_threads_(int* num_threads){
 
 }
+
+int openblas_set_num_threads_local(int num_threads){
+	return 1;
+}
 #endif