From e783a732771ea1129ea4cdfec9517bf71b916069 Mon Sep 17 00:00:00 2001
From: TengLu <teng.lu@intel.com>
Date: Tue, 5 Nov 2019 09:19:06 +0800
Subject: [PATCH] Enable inter env var for MKL-DNN backend.

---
 tensorflow/core/common_runtime/process_util.cc | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/common_runtime/process_util.cc b/tensorflow/core/common_runtime/process_util.cc
index 9794969827d..d4674cad64e 100644
--- a/tensorflow/core/common_runtime/process_util.cc
+++ b/tensorflow/core/common_runtime/process_util.cc
@@ -34,10 +34,15 @@ namespace tensorflow {
 
 namespace {
 
+// Use environment setting if specified (init once)
+int32 GetEnvNumInterOpThreads() {
+  static int32 env_num_threads = NumInterOpThreadsFromEnvironment();
+  return env_num_threads;
+}
+
 int32 DefaultNumInterOpThreads() {
 #ifndef __ANDROID__
-  // Use environment setting if specified (init once)
-  static int env_num_threads = NumInterOpThreadsFromEnvironment();
+  int32 env_num_threads = GetEnvNumInterOpThreads();
   if (env_num_threads > 0) {
     return env_num_threads;
   }
@@ -121,6 +126,9 @@ int32 DefaultNumIntraOpThreads() {
 int32 NumInterOpThreadsFromSessionOptions(const SessionOptions& options) {
   const int32 inter_op = options.config.inter_op_parallelism_threads();
   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
   if (!DisableMKL()) {
     // MKL library executes ops in parallel using OMP threads.
@@ -163,7 +171,7 @@ void SchedClosure(std::function<void()> closure) {
   uint64 id = tracing::GetUniqueArg();
   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);
     closure();
   });