torch.set_num_threads sets MKL option too (#4949)

soumith · web-flow · commit 2aaeec0db0be · 2018-01-31T11:59:15.000-05:00
diff --git a/aten/src/TH/THGeneral.c b/aten/src/TH/THGeneral.c
@@ -21,6 +21,11 @@
 #include <malloc/malloc.h>
 #endif
 
+#ifdef TH_BLAS_MKL
+extern void mkl_set_num_threads(int);
+extern int mkl_get_max_threads(void);
+#endif
+
 /* Torch Error Handling */
 static void defaultErrorHandlerFunction(const char *msg, void *data)
 {
@@ -302,6 +307,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,10 +331,6 @@ 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)
