From b2edec9720abff1519460b141b64a11f8fdf70cc Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Thu, 1 Mar 2018 14:23:05 +0000 Subject: [PATCH] [Minor] Support threads for MKL blas --- contrib/torch/torch7/lib/TH/THGeneral.c | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/contrib/torch/torch7/lib/TH/THGeneral.c b/contrib/torch/torch7/lib/TH/THGeneral.c index 50dba205f..b9491c9f3 100644 --- a/contrib/torch/torch7/lib/TH/THGeneral.c +++ b/contrib/torch/torch7/lib/TH/THGeneral.c @@ -328,6 +328,11 @@ void THSetNumThreads(int num_threads) extern void openblas_set_num_threads(int); openblas_set_num_threads(num_threads); #endif +#ifdef TH_BLAS_MKL + extern void mkl_set_num_threads(int); + mkl_set_num_threads(num_threads); + +#endif } int THGetNumThreads(void) @@ -342,7 +347,12 @@ int THGetNumThreads(void) bl_threads = openblas_get_num_threads(); nthreads = nthreads > bl_threads ? bl_threads : nthreads; #endif - +#ifdef TH_BLAS_MKL + int bl_threads = 1; + extern int mkl_get_max_threads(void); + bl_threads = mkl_get_max_threads(); + nthreads = nthreads > bl_threads ? bl_threads : nthreads; +#endif return nthreads; } -- 2.39.5