experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Enable inter env var for MKL-DNN backend.

Browse Source
This commit is contained in:
TengLu 2019-11-05 09:19:06 +08:00
parent e6f23d3e45
commit e783a73277
1 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tensorflow/core/common_runtime/process_util.cc
View File

@ -34,10 +34,15 @@ namespace tensorflow {
namespace { namespace {
// Use environment setting if specified (init once)
int32 GetEnvNumInterOpThreads() {
static int32 env_num_threads = NumInterOpThreadsFromEnvironment();
return env_num_threads;
}
int32 DefaultNumInterOpThreads() { int32 DefaultNumInterOpThreads() {
#ifndef __ANDROID__ #ifndef __ANDROID__
// Use environment setting if specified (init once) int32 env_num_threads = GetEnvNumInterOpThreads();
static int env_num_threads = NumInterOpThreadsFromEnvironment();
if (env_num_threads > 0) { if (env_num_threads > 0) {
return env_num_threads; return env_num_threads;
} }
@ -121,6 +126,9 @@ int32 DefaultNumIntraOpThreads() {
int32 NumInterOpThreadsFromSessionOptions(const SessionOptions& options) { int32 NumInterOpThreadsFromSessionOptions(const SessionOptions& options) {
const int32 inter_op = options.config.inter_op_parallelism_threads(); const int32 inter_op = options.config.inter_op_parallelism_threads();
if (inter_op > 0) return inter_op; if (inter_op > 0) return inter_op;
const int32 env_inter_op = GetEnvNumInterOpThreads();
if (env_inter_op > 0) return env_inter_op;
#ifdef INTEL_MKL #ifdef INTEL_MKL
if (!DisableMKL()) { if (!DisableMKL()) {
// MKL library executes ops in parallel using OMP threads. // MKL library executes ops in parallel using OMP threads.
@ -163,7 +171,7 @@ void SchedClosure(std::function<void()> closure) {
uint64 id = tracing::GetUniqueArg(); uint64 id = tracing::GetUniqueArg();
tracing::RecordEvent(tracing::EventCategory::kScheduleClosure, id); tracing::RecordEvent(tracing::EventCategory::kScheduleClosure, id);
Env::Default()->SchedClosure([id, closure = std::move(closure)]() { Env::Default()->SchedClosure([ id, closure = std::move(closure) ]() {
tracing::ScopedRegion region(tracing::EventCategory::kRunClosure, id); tracing::ScopedRegion region(tracing::EventCategory::kRunClosure, id);
closure(); closure();
}); });