openblas_threads_callback mechanism for changing threads backend

cblas.h

@@ -25,6 +25,11 @@ char* openblas_get_config(void);
 /*Get the CPU corename on runtime.*/
 char* openblas_get_corename(void);
 
+/*Set the threading backend to a custom callback.*/
+typedef void (*openblas_dojob_callback)(int thread_num, void *jobdata, void *dojob_data);
+typedef void (*openblas_threads_callback)(void *callback_data, openblas_dojob_callback dojob, int numjobs, size_t jobdata_elsize, void *jobdata, void *dojob_data);
+void openblas_set_threads_callback(openblas_threads_callback callback, void *callback_data);
+
 /* Get the parallelization type which is used by OpenBLAS */
 int openblas_get_parallel(void);
 /* OpenBLAS is compiled for sequential use  */

common_interface.h

@@ -47,6 +47,12 @@ int    BLASFUNC(xerbla)(char *, blasint *info, blasint);
 
 void    openblas_set_num_threads_(int *);
 
+typedef void (*openblas_dojob_callback)(int thread_num, void *jobdata, void *dojob_data);
+typedef void (*openblas_threads_callback)(void *callback_data, openblas_dojob_callback dojob, int numjobs, size_t jobdata_elsize, void *jobdata, void *dojob_data);
+void openblas_set_threads_callback(openblas_threads_callback callback, void *callback_data);
+extern openblas_threads_callback openblas_threads_callback_;
+extern void *openblas_threads_callback_data_;
+
 FLOATRET  BLASFUNC(sdot)  (blasint *, float  *, blasint *, float  *, blasint *);
 FLOATRET  BLASFUNC(sdsdot)(blasint *, float  *,        float  *, blasint *, float  *, blasint *);
 

driver/others/CMakeLists.txt

@@ -39,6 +39,7 @@ set(COMMON_SOURCES
   openblas_env.c
   openblas_get_num_procs.c
   openblas_get_num_threads.c
+  blas_server_callback.c
 )
 
 # these need to have NAME/CNAME set, so use GenerateNamedObjects, but don't use standard name mangling

driver/others/Makefile

@@ -1,7 +1,7 @@
 TOPDIR	= ../..
 include ../../Makefile.system
 
-COMMONOBJS	 = memory.$(SUFFIX) xerbla.$(SUFFIX) c_abs.$(SUFFIX) z_abs.$(SUFFIX) openblas_set_num_threads.$(SUFFIX) openblas_get_num_threads.$(SUFFIX) openblas_get_num_procs.$(SUFFIX) openblas_get_config.$(SUFFIX) openblas_get_parallel.$(SUFFIX) openblas_error_handle.$(SUFFIX) openblas_env.$(SUFFIX)
+COMMONOBJS	 = memory.$(SUFFIX) xerbla.$(SUFFIX) c_abs.$(SUFFIX) z_abs.$(SUFFIX) openblas_set_num_threads.$(SUFFIX) openblas_get_num_threads.$(SUFFIX) openblas_get_num_procs.$(SUFFIX) openblas_get_config.$(SUFFIX) openblas_get_parallel.$(SUFFIX) openblas_error_handle.$(SUFFIX) openblas_env.$(SUFFIX) blas_server_callback.$(SUFFIX)
 
 #COMMONOBJS	+= slamch.$(SUFFIX) slamc3.$(SUFFIX) dlamch.$(SUFFIX)  dlamc3.$(SUFFIX)
 

driver/others/blas_server_callback.c

@@ -0,0 +1,13 @@
+#include "common.h"
+
+/* global variable to change threading backend from openblas-managed to caller-managed */
+openblas_threads_callback openblas_threads_callback_ = 0;
+void *openblas_threads_callback_data_ = 0;
+
+/* non-threadsafe function should be called before any other
+   openblas function to change how threads are managed */
+void openblas_set_threads_callback(openblas_threads_callback callback, void *callback_data)
+{
+  openblas_threads_callback_ = callback;
+  openblas_threads_callback_data_ = callback_data;
+}

driver/others/blas_server_omp.c

@@ -75,6 +75,7 @@ void goto_set_num_threads(int num_threads) {
 
   blas_cpu_number  = num_threads;
 
+  if (!openblas_threads_callback_)
     omp_set_num_threads(blas_cpu_number);
 
   //adjust buffer for each thread
@@ -222,10 +223,9 @@ static void legacy_exec(void *func, int mode, blas_arg_t *args, void *sb){
       }
 }
 
-static void exec_threads(blas_queue_t *queue, int buf_index){
+static void exec_threads(int thread_num, blas_queue_t *queue, int *buf_index){
-
   void *buffer, *sa, *sb;
-  int pos=0, release_flag=0;
+  int release_flag=0;
 
   buffer = NULL;
   sa = queue -> sa;
@@ -238,8 +238,7 @@ static void exec_threads(blas_queue_t *queue, int buf_index){
 
   if ((sa == NULL) && (sb == NULL) && ((queue -> mode & BLAS_PTHREAD) == 0)) {
 
-    pos = omp_get_thread_num();
+    buffer = blas_thread_buffer[*buf_index][thread_num];
-    buffer = blas_thread_buffer[buf_index][pos];
 
     //fallback
     if(buffer==NULL) {
@@ -335,14 +334,25 @@ int exec_blas(BLASLONG num, blas_queue_t *queue){
       break;
   }
 
+  if (openblas_threads_callback_) {
+#ifndef USE_SIMPLE_THREADED_LEVEL3
+    for (i = 0; i < num; i ++)
+      queue[i].position = i;
+#endif
+    openblas_threads_callback_(openblas_threads_callback_data_, (openblas_dojob_callback) exec_threads, num, sizeof(blas_queue_t), (void*) queue, (void*) &buf_index);
+    return;
+  }
+
 #pragma omp parallel for schedule(OMP_SCHED)
+
+#pragma omp parallel for schedule(static)
   for (i = 0; i < num; i ++) {
 
 #ifndef USE_SIMPLE_THREADED_LEVEL3
     queue[i].position = i;
 #endif
 
-    exec_threads(&queue[i], buf_index);
+    exec_threads(omp_get_thread_num(), &queue[i], &buf_index);
   }
 
 #if __STDC_VERSION__ >= 201112L