Caffe2 - C++ API
A deep learning, cross platform ML framework
init_omp.cc
1 
17 #include <stdlib.h>
18 
19 #include "caffe2/core/common.h"
20 
21 #ifdef _OPENMP
22 #include "caffe2/core/common_omp.h"
23 #endif // _OPENMP
24 
25 #ifdef CAFFE2_USE_MKL
26 #include <mkl.h>
27 #endif // CAFFE2_USE_MKL
28 
29 #include "caffe2/core/init.h"
30 
31 CAFFE2_DEFINE_int(
32  caffe2_omp_num_threads, 0,
33  "The number of openmp threads. 0 to use default value. "
34  "Does not have effect if OpenMP is disabled.");
35 CAFFE2_DEFINE_int(
36  caffe2_mkl_num_threads,
37  0,
38  "The number of mkl threads. 0 to use default value. If set, "
39  "this overrides the caffe2_omp_num_threads flag if both are set. "
40  "Does not have effect if MKL is not used.");
41 
42 namespace caffe2 {
43 
44 #ifdef _OPENMP
45 bool Caffe2SetOpenMPThreads(int*, char***) {
46  if (!getenv("OMP_NUM_THREADS")) {
47  // OMP_NUM_THREADS not passed explicitly, so *disable* OMP by
48  // default. The user can use the CLI flag to override.
49  VLOG(1) << "OMP_NUM_THREADS not passed, defaulting to 1 thread";
50  omp_set_num_threads(1);
51  }
52 
53  if (FLAGS_caffe2_omp_num_threads > 0) {
54  VLOG(1) << "Setting omp_num_threads to " << FLAGS_caffe2_omp_num_threads;
55  omp_set_num_threads(FLAGS_caffe2_omp_num_threads);
56  }
57  VLOG(1) << "Caffe2 running with " << omp_get_max_threads() << " OMP threads";
58  return true;
59 }
60 REGISTER_CAFFE2_INIT_FUNCTION(Caffe2SetOpenMPThreads,
61  &Caffe2SetOpenMPThreads,
62  "Set OpenMP threads.");
63 #endif // _OPENMP
64 
65 #ifdef CAFFE2_USE_MKL
66 bool Caffe2SetMKLThreads(int*, char***) {
67  if (!getenv("MKL_NUM_THREADS")) {
68  VLOG(1) << "MKL_NUM_THREADS not passed, defaulting to 1 thread";
69  mkl_set_num_threads(1);
70  }
71 
72  // If caffe2_omp_num_threads is set, we use that for MKL as well.
73  if (FLAGS_caffe2_omp_num_threads > 0) {
74  VLOG(1) << "Setting mkl_num_threads to " << FLAGS_caffe2_omp_num_threads
75  << " as inherited from omp_num_threads.";
76  mkl_set_num_threads(FLAGS_caffe2_omp_num_threads);
77  }
78 
79  // Override omp_num_threads if mkl_num_threads is set.
80  if (FLAGS_caffe2_mkl_num_threads > 0) {
81  VLOG(1) << "Setting mkl_num_threads to " << FLAGS_caffe2_mkl_num_threads;
82  mkl_set_num_threads(FLAGS_caffe2_mkl_num_threads);
83  }
84  VLOG(1) << "Caffe2 running with " << mkl_get_max_threads() << " MKL threads";
85  return true;
86 }
87 REGISTER_CAFFE2_INIT_FUNCTION(
88  Caffe2SetMKLThreads,
89  &Caffe2SetMKLThreads,
90  "Set MKL threads.");
91 #endif // CAFFE2_USE_MKL
92 
93 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.