From 3c3a11a0dc780498a7c890be90b9df922b426d90 Mon Sep 17 00:00:00 2001
From: tensor-tang <jian.j.tang@intel.com>
Date: Fri, 4 Aug 2017 13:50:41 +0800
Subject: [PATCH] add use_mkldnn flag

---
 paddle/gserver/layers/MkldnnLayer.h    |  4 +++-
 paddle/trainer/TrainerConfigHelper.cpp |  2 ++
 paddle/utils/Flags.cpp                 |  7 +++++++
 paddle/utils/Flags.h                   |  1 +
 python/paddle/trainer/config_parser.py | 24 +++++++++++++++++++++---
 5 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/paddle/gserver/layers/MkldnnLayer.h b/paddle/gserver/layers/MkldnnLayer.h
index e9bab68b07..7e6d88b273 100644
--- a/paddle/gserver/layers/MkldnnLayer.h
+++ b/paddle/gserver/layers/MkldnnLayer.h
@@ -35,8 +35,10 @@ public:
 
   virtual bool init(const LayerMap& layerMap,
                     const ParameterMap& parameterMap) {
-    return Layer::init(layerMap, parameterMap);
+    CHECK(FLAGS_use_mkldnn) << "MkldnnLayers only support use_mkldnn."
+                            << "Please set WITH_MKLDNN=ON";
     // TODO(TJ): deivecId
+    return Layer::init(layerMap, parameterMap);
   }
 
   void resetOutput(size_t height, size_t width) { ; }
diff --git a/paddle/trainer/TrainerConfigHelper.cpp b/paddle/trainer/TrainerConfigHelper.cpp
index 133e2be104..a0a365aa0b 100644
--- a/paddle/trainer/TrainerConfigHelper.cpp
+++ b/paddle/trainer/TrainerConfigHelper.cpp
@@ -28,6 +28,7 @@ DECLARE_bool(with_cost);
 DECLARE_bool(with_gpu);
 DECLARE_bool(parallel_nn);
 DECLARE_string(config_args);
+DECLARE_bool(use_mkldnn);
 
 const char *kConfigParserModuleName = "paddle.trainer.config_parser";
 const char *kConfigParserFuncName = "parse_config_and_serialize";
@@ -44,6 +45,7 @@ TrainerConfigHelper::TrainerConfigHelper(const std::string &configFilePath)
   configArgs << "trainer_id=" << FLAGS_trainer_id << ",local=" << FLAGS_local
              << ",with_cost=" << FLAGS_with_cost << ",use_gpu=" << FLAGS_use_gpu
              << ",parallel_nn=" << FLAGS_parallel_nn
+             << ",use_mkldnn=" << FLAGS_use_mkldnn
              << ",cudnn_version=" << hl_get_cudnn_lib_version();
   if (!FLAGS_config_args.empty()) {
     configArgs << "," << FLAGS_config_args;
diff --git a/paddle/utils/Flags.cpp b/paddle/utils/Flags.cpp
index 320f671ed9..ab1c181c62 100644
--- a/paddle/utils/Flags.cpp
+++ b/paddle/utils/Flags.cpp
@@ -20,6 +20,13 @@ DEFINE_bool(use_gpu, false, "Only support CPU training");
 DEFINE_bool(use_gpu, true, "Whether to use GPU for training");
 #endif
 
+#ifdef PADDLE_USE_MKLDNN
+// TODO(TJ): change to true when MKLDNN layers support multi-inputs
+DEFINE_bool(use_mkldnn, false, "Default still keep use CPU training");
+#else
+DEFINE_bool(use_mkldnn, false, "Only support CPU training");
+#endif
+
 DEFINE_bool(parallel_nn,
             false,
             "Whether to use multi-threads to calculate one neural network."
diff --git a/paddle/utils/Flags.h b/paddle/utils/Flags.h
index dc4faef833..1832bb515e 100644
--- a/paddle/utils/Flags.h
+++ b/paddle/utils/Flags.h
@@ -40,3 +40,4 @@ DECLARE_bool(show_layer_stat);
 DECLARE_string(predict_file);
 DECLARE_bool(prev_batch_state);
 DECLARE_string(init_model_path);
+DECLARE_bool(use_mkldnn);
diff --git a/python/paddle/trainer/config_parser.py b/python/paddle/trainer/config_parser.py
index 9ea69fc5e5..ae39abc081 100644
--- a/python/paddle/trainer/config_parser.py
+++ b/python/paddle/trainer/config_parser.py
@@ -1604,6 +1604,8 @@ class MultiClassCrossEntropySelfNormCostLayer(LayerBase):
 
 @config_layer('fc')
 class FCLayer(LayerBase):
+    layer_type = 'fc'
+
     def __init__(self,
                  name,
                  size,
@@ -1611,14 +1613,25 @@ class FCLayer(LayerBase):
                  bias=True,
                  error_clipping_threshold=None,
                  **xargs):
-        super(FCLayer, self).__init__(name, 'fc', size, inputs=inputs, **xargs)
+        use_mkldnn = bool(int(g_command_config_args.get("use_mkldnn", 0)))
+        if use_mkldnn:
+            self.layer_type = 'mkldnn_fc'
+            config_assert(
+                len(inputs) == 1,
+                "MkldnnFCLayer support one and only one input!")
+        super(FCLayer, self).__init__(
+            name, self.layer_type, size, inputs=inputs, **xargs)
         for input_index in xrange(len(self.inputs)):
             input_layer = self.get_input_layer(input_index)
             psize = self.config.size * input_layer.size
-            dims = [input_layer.size, self.config.size]
             format = self.inputs[input_index].format
             sparse = format == "csr" or format == "csc"
-
+            if use_mkldnn:
+                dims = [self.config.size, input_layer.size]
+                config_assert(not sparse,
+                              "MkldnnFCLayer do not support sparse format yet")
+            else:
+                dims = [input_layer.size, self.config.size]
             if sparse:
                 psize = self.inputs[input_index].nnz
             else:
@@ -1631,6 +1644,11 @@ class FCLayer(LayerBase):
             self.config.error_clipping_threshold = error_clipping_threshold
 
 
+@config_layer('mkldnn_fc')
+class MkldnnFcLayer(FCLayer):
+    layer_type = 'mkldnn_fc'
+
+
 @config_layer('selective_fc')
 class SelectiveFCLayer(LayerBase):
     def __init__(self,
-- 
GitLab