From 8376ff25e65545d0da9cb91dfff587af6a17e914 Mon Sep 17 00:00:00 2001
From: Xin Qiu <qiuxin2012@users.noreply.github.com>
Date: Wed, 11 Sep 2019 11:01:04 +0800
Subject: [PATCH] add multi thread sample to minibatch (#1589)

* add mt sample to minibatch

* revert log level

* delete some debug code

* add comments

* fix style check

* update pom
---
 .../bigdl/dllib/utils/NNContext.scala         | 58 ++++++++++---------
 1 file changed, 31 insertions(+), 27 deletions(-)

diff --git a/scala/dllib/src/main/scala/com/intel/analytics/bigdl/dllib/utils/NNContext.scala b/scala/dllib/src/main/scala/com/intel/analytics/bigdl/dllib/utils/NNContext.scala
index 9d7f8d9c49e..7c0f41bef55 100644
--- a/scala/dllib/src/main/scala/com/intel/analytics/bigdl/dllib/utils/NNContext.scala
+++ b/scala/dllib/src/main/scala/com/intel/analytics/bigdl/dllib/utils/NNContext.scala
@@ -209,35 +209,39 @@ object NNContext {
   private[zoo] def initConf(zooConf: SparkConf) : Unit = {
     // check env and set spark conf
     // Set default value
-    var kmpAffinity = "granularity=fine,compact,1,0"
-    var kmpBlockTime = "0"
-    var kmpSettings = "1"
-    var ompNumThreads = "1"
-    // Set value with env
-    if (env.contains("KMP_AFFINITY")) {
-      kmpAffinity = env("KMP_AFFINITY")
-    }
-    if (env.contains("KMP_BLOCKTIME")) {
-      kmpBlockTime = env("KMP_BLOCKTIME")
-    }
-    if (env.contains("KMP_SETTINGS")) {
-      kmpSettings = env("KMP_SETTINGS")
-    }
-    if (env.contains("OMP_NUM_THREADS")) {
-      ompNumThreads = env("OMP_NUM_THREADS")
-    } else if (env.contains("ZOO_NUM_MKLTHREADS")) {
-      if (env("ZOO_NUM_MKLTHREADS").equalsIgnoreCase("all")) {
-        val cores = Runtime.getRuntime.availableProcessors()
-        ompNumThreads = cores.toString
-      } else {
-        ompNumThreads = env("ZOO_NUM_MKLTHREADS")
+    // We should skip this env, when engineType is mkldnn.
+    if (System.getProperty("bigdl.engineType", "mklblas")
+      .toLowerCase() == "mklblas") {
+      var kmpAffinity = "granularity=fine,compact,1,0"
+      var kmpBlockTime = "0"
+      var kmpSettings = "1"
+      var ompNumThreads = "1"
+      // Set value with env
+      if (env.contains("KMP_AFFINITY")) {
+        kmpAffinity = env("KMP_AFFINITY")
+      }
+      if (env.contains("KMP_BLOCKTIME")) {
+        kmpBlockTime = env("KMP_BLOCKTIME")
+      }
+      if (env.contains("KMP_SETTINGS")) {
+        kmpSettings = env("KMP_SETTINGS")
+      }
+      if (env.contains("OMP_NUM_THREADS")) {
+        ompNumThreads = env("OMP_NUM_THREADS")
+      } else if (env.contains("ZOO_NUM_MKLTHREADS")) {
+        if (env("ZOO_NUM_MKLTHREADS").equalsIgnoreCase("all")) {
+          val cores = Runtime.getRuntime.availableProcessors()
+          ompNumThreads = cores.toString
+        } else {
+          ompNumThreads = env("ZOO_NUM_MKLTHREADS")
+        }
       }
+      // Set Spark Conf
+      zooConf.setExecutorEnv("KMP_AFFINITY", kmpAffinity)
+      zooConf.setExecutorEnv("KMP_BLOCKTIME", kmpBlockTime)
+      zooConf.setExecutorEnv("KMP_SETTINGS", kmpSettings)
+      zooConf.setExecutorEnv("OMP_NUM_THREADS", ompNumThreads)
     }
-    // Set Spark Conf
-    zooConf.setExecutorEnv("KMP_AFFINITY", kmpAffinity)
-    zooConf.setExecutorEnv("KMP_BLOCKTIME", kmpBlockTime)
-    zooConf.setExecutorEnv("KMP_SETTINGS", kmpSettings)
-    zooConf.setExecutorEnv("OMP_NUM_THREADS", ompNumThreads)
 
   }