make torch.set_num_threads also set MKL threads (take 2) (#5002)

soumith · web-flow · commit bf3655a10c8e · 2018-02-02T09:24:54.000-05:00
* torch.set_num_threads sets MKL option too

* fix to use C prototype instead of fortran
diff --git a/aten/src/TH/THGeneral.c b/aten/src/TH/THGeneral.c
@@ -21,6 +21,13 @@
 #include <malloc/malloc.h>
 #endif
 
+#ifdef TH_BLAS_MKL
+// this is the C prototype, while mkl_set_num_threads is the fortran prototype
+extern void MKL_Set_Num_Threads(int);
+// this is the C prototype, while mkl_get_max_threads is the fortran prototype
+extern int  MKL_Get_Max_Threads(void);
+#endif
+
 /* Torch Error Handling */
 static void defaultErrorHandlerFunction(const char *msg, void *data)
 {
@@ -302,6 +309,10 @@ void THSetNumThreads(int num_threads)
 #ifdef _OPENMP
   omp_set_num_threads(num_threads);
 #endif
+#ifdef TH_BLAS_MKL
+  MKL_Set_Num_Threads(num_threads);
+#endif
+
 }
 
 int THGetNumThreads(void)
@@ -322,18 +333,14 @@ int THGetNumCores(void)
 #endif
 }
 
-#ifdef TH_BLAS_MKL
-extern int mkl_get_max_threads(void);
-#endif
-
 TH_API void THInferNumThreads(void)
 {
 #if defined(_OPENMP) && defined(TH_BLAS_MKL)
   // If we are using MKL an OpenMP make sure the number of threads match.
   // Otherwise, MKL and our OpenMP-enabled functions will keep changing the
   // size of the OpenMP thread pool, resulting in worse performance (and memory
   // leaks in GCC 5.4)
-  omp_set_num_threads(mkl_get_max_threads());
+  omp_set_num_threads(MKL_Get_Max_Threads());
 #endif
 }
 
