add multi thread sample to minibatch (intel-analytics#1589)

* add mt sample to minibatch * revert log level * delete some debug code * add comments * fix style check * update pom
dding3 · Sep 11, 2019 · bfaa989 · bfaa989
1 parent 24e478b
commit bfaa989
Showing 1 changed file with 31 additions and 27 deletions.
diff --git a/NNContext.scala b/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)
 
   }