mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
7d5f7ed270
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/12714 This is a short change to enable c10 namespace in caffe2. We did not enable it before due to gflags global variable confusion, but it should have been mostly cleaned now. Right now, the plan on record is that namespace caffe2 and namespace aten will fully be supersets of namespace c10. Most of the diff is codemod, and only two places of non-codemod is in caffe2/core/common.h, where ``` using namespace c10; ``` is added, and in Flags.h, where instead of creating aliasing variables in c10 namespace, we directly put it in the global namespace to match gflags (and same behavior if gflags is not being built with). Reviewed By: dzhulgakov Differential Revision: D10390486 fbshipit-source-id: 5e2df730e28e29a052f513bddc558d9f78a23b9b
79 lines
2.4 KiB
C++
79 lines
2.4 KiB
C++
#include <stdlib.h>
|
|
|
|
#include "caffe2/core/common.h"
|
|
|
|
#ifdef _OPENMP
|
|
#include "caffe2/core/common_omp.h"
|
|
#endif // _OPENMP
|
|
|
|
#ifdef CAFFE2_USE_MKL
|
|
#include <mkl.h>
|
|
#endif // CAFFE2_USE_MKL
|
|
|
|
#include "caffe2/core/init.h"
|
|
|
|
C10_DEFINE_int(
|
|
caffe2_omp_num_threads,
|
|
0,
|
|
"The number of openmp threads. 0 to use default value. "
|
|
"Does not have effect if OpenMP is disabled.");
|
|
C10_DEFINE_int(
|
|
caffe2_mkl_num_threads,
|
|
0,
|
|
"The number of mkl threads. 0 to use default value. If set, "
|
|
"this overrides the caffe2_omp_num_threads flag if both are set. "
|
|
"Does not have effect if MKL is not used.");
|
|
|
|
namespace caffe2 {
|
|
|
|
#ifdef _OPENMP
|
|
bool Caffe2SetOpenMPThreads(int*, char***) {
|
|
if (!getenv("OMP_NUM_THREADS")) {
|
|
// OMP_NUM_THREADS not passed explicitly, so *disable* OMP by
|
|
// default. The user can use the CLI flag to override.
|
|
VLOG(1) << "OMP_NUM_THREADS not passed, defaulting to 1 thread";
|
|
omp_set_num_threads(1);
|
|
}
|
|
|
|
if (FLAGS_caffe2_omp_num_threads > 0) {
|
|
VLOG(1) << "Setting omp_num_threads to " << FLAGS_caffe2_omp_num_threads;
|
|
omp_set_num_threads(FLAGS_caffe2_omp_num_threads);
|
|
}
|
|
VLOG(1) << "Caffe2 running with " << omp_get_max_threads() << " OMP threads";
|
|
return true;
|
|
}
|
|
REGISTER_CAFFE2_INIT_FUNCTION(Caffe2SetOpenMPThreads,
|
|
&Caffe2SetOpenMPThreads,
|
|
"Set OpenMP threads.");
|
|
#endif // _OPENMP
|
|
|
|
#ifdef CAFFE2_USE_MKL
|
|
bool Caffe2SetMKLThreads(int*, char***) {
|
|
if (!getenv("MKL_NUM_THREADS")) {
|
|
VLOG(1) << "MKL_NUM_THREADS not passed, defaulting to 1 thread";
|
|
mkl_set_num_threads(1);
|
|
}
|
|
|
|
// If caffe2_omp_num_threads is set, we use that for MKL as well.
|
|
if (FLAGS_caffe2_omp_num_threads > 0) {
|
|
VLOG(1) << "Setting mkl_num_threads to " << FLAGS_caffe2_omp_num_threads
|
|
<< " as inherited from omp_num_threads.";
|
|
mkl_set_num_threads(FLAGS_caffe2_omp_num_threads);
|
|
}
|
|
|
|
// Override omp_num_threads if mkl_num_threads is set.
|
|
if (FLAGS_caffe2_mkl_num_threads > 0) {
|
|
VLOG(1) << "Setting mkl_num_threads to " << FLAGS_caffe2_mkl_num_threads;
|
|
mkl_set_num_threads(FLAGS_caffe2_mkl_num_threads);
|
|
}
|
|
VLOG(1) << "Caffe2 running with " << mkl_get_max_threads() << " MKL threads";
|
|
return true;
|
|
}
|
|
REGISTER_CAFFE2_INIT_FUNCTION(
|
|
Caffe2SetMKLThreads,
|
|
&Caffe2SetMKLThreads,
|
|
"Set MKL threads.");
|
|
#endif // CAFFE2_USE_MKL
|
|
|
|
} // namespace caffe2
|