3 #include "caffe2/core/common.h" 6 #include "caffe2/core/common_omp.h" 11 #endif // CAFFE2_USE_MKL 13 #include "caffe2/core/init.h" 16 caffe2_omp_num_threads, 0,
17 "The number of openmp threads. 0 to use default value. " 18 "Does not have effect if OpenMP is disabled.");
20 caffe2_mkl_num_threads,
22 "The number of mkl threads. 0 to use default value. If set, " 23 "this overrides the caffe2_omp_num_threads flag if both are set. " 24 "Does not have effect if MKL is not used.");
29 bool Caffe2SetOpenMPThreads(
int*,
char***) {
30 if (!getenv(
"OMP_NUM_THREADS")) {
33 VLOG(1) <<
"OMP_NUM_THREADS not passed, defaulting to 1 thread";
34 omp_set_num_threads(1);
37 if (FLAGS_caffe2_omp_num_threads > 0) {
38 VLOG(1) <<
"Setting omp_num_threads to " << FLAGS_caffe2_omp_num_threads;
39 omp_set_num_threads(FLAGS_caffe2_omp_num_threads);
41 VLOG(1) <<
"Caffe2 running with " << omp_get_max_threads() <<
" OMP threads";
44 REGISTER_CAFFE2_INIT_FUNCTION(Caffe2SetOpenMPThreads,
45 &Caffe2SetOpenMPThreads,
46 "Set OpenMP threads.");
50 bool Caffe2SetMKLThreads(
int*,
char***) {
51 if (!getenv(
"MKL_NUM_THREADS")) {
52 VLOG(1) <<
"MKL_NUM_THREADS not passed, defaulting to 1 thread";
53 mkl_set_num_threads(1);
57 if (FLAGS_caffe2_omp_num_threads > 0) {
58 VLOG(1) <<
"Setting mkl_num_threads to " << FLAGS_caffe2_omp_num_threads
59 <<
" as inherited from omp_num_threads.";
60 mkl_set_num_threads(FLAGS_caffe2_omp_num_threads);
64 if (FLAGS_caffe2_mkl_num_threads > 0) {
65 VLOG(1) <<
"Setting mkl_num_threads to " << FLAGS_caffe2_mkl_num_threads;
66 mkl_set_num_threads(FLAGS_caffe2_mkl_num_threads);
68 VLOG(1) <<
"Caffe2 running with " << mkl_get_max_threads() <<
" MKL threads";
71 REGISTER_CAFFE2_INIT_FUNCTION(
75 #endif // CAFFE2_USE_MKL A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...