Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into add-async-ssa-graph-executor

.pre-commit-config.yaml

@@ -42,12 +42,6 @@ repos:
         entry: bash ./tools/codestyle/pylint_pre_commit.hook
         language: system
         files: \.(py)$
--   repo: https://github.com/PaddlePaddle/pre-commit-golang
-    sha: 8337620115c25ff8333f1b1a493bd031049bd7c0
-    hooks:
-    -   id: go-fmt
-        types:
-        - go
 -   repo: local
     hooks:
     -   id: copyright_checker

cmake/external/mklml.cmake

@@ -39,8 +39,10 @@ IF(WIN32)
     SET(MKLML_IOMP_LIB            ${MKLML_LIB_DIR}/libiomp5md.lib)
     SET(MKLML_SHARED_LIB          ${MKLML_LIB_DIR}/mklml.dll)
     SET(MKLML_SHARED_IOMP_LIB     ${MKLML_LIB_DIR}/libiomp5md.dll)
-ELSE()  
-    SET(MKLML_VER "mklml_lnx_${TIME_VERSION}" CACHE STRING "" FORCE)
+ELSE()
+    #TODO(intel-huying):
+    #  Now enable Erf function in mklml library temporarily, it will be updated as offical version later.
+    SET(MKLML_VER "Glibc225_vsErf_mklml_lnx_${TIME_VERSION}" CACHE STRING "" FORCE)
     SET(MKLML_URL "http://paddlepaddledeps.cdn.bcebos.com/${MKLML_VER}.tgz" CACHE STRING "" FORCE)
     SET(MKLML_LIB                 ${MKLML_LIB_DIR}/libmklml_intel.so)
     SET(MKLML_IOMP_LIB            ${MKLML_LIB_DIR}/libiomp5.so)

cmake/external/ngraph.cmake

@@ -37,7 +37,7 @@ INCLUDE(GNUInstallDirs)
 INCLUDE(ExternalProject)
 
 SET(NGRAPH_PROJECT         "extern_ngraph")
-SET(NGRAPH_GIT_TAG         "20bd8bbc79ae3a81c57313846a2be7313e5d1dab")
+SET(NGRAPH_GIT_TAG         "a444f7a959b7d87f2c117c9b57a4c387759e481e")
 SET(NGRAPH_SOURCES_DIR     ${THIRD_PARTY_PATH}/ngraph)
 SET(NGRAPH_INSTALL_DIR     ${THIRD_PARTY_PATH}/install/ngraph)
 SET(NGRAPH_INC_DIR         ${NGRAPH_INSTALL_DIR}/include)
@@ -69,7 +69,7 @@ ExternalProject_Add(
     CMAKE_ARGS          -DNGRAPH_DEX_ONLY=TRUE
     CMAKE_ARGS          -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
     CMAKE_ARGS          -DMKLDNN_INCLUDE_DIR=${MKLDNN_INC_DIR}
-    CMAKE_ARGS          -DMKLDNN_LIB_DIR=${MKLDNN_INSTALL_DIR}/lib
+    CMAKE_ARGS          -DMKLDNN_LIB_DIR=${MKLDNN_INSTALL_DIR}/${CMAKE_INSTALL_LIBDIR}
     CMAKE_ARGS          -DMKLML_LIB_DIR=${MKLML_INSTALL_DIR}/lib
 )
 

cmake/operators.cmake

@@ -153,7 +153,11 @@ function(op_library TARGET)
     # pybind USE_OP_DEVICE_KERNEL for CUDNN
     list(LENGTH cudnn_cu_cc_srcs cudnn_cu_cc_srcs_len)
     if (WITH_GPU AND ${cudnn_cu_cc_srcs_len} GREATER 0)
+      if(${TARGET} STREQUAL "activation")
+        file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL(relu, CUDNN);\n")
+      else()
         file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL(${TARGET}, CUDNN);\n")
+      endif()
     endif()
 
     # pybind USE_OP_DEVICE_KERNEL for MIOPEN
@@ -168,6 +172,9 @@ function(op_library TARGET)
         file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL(relu, MKLDNN);\n")
       elseif(${MKLDNN_FILE} STREQUAL "conv_mkldnn_op")
         file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL_WITH_CUSTOM_TYPE(conv2d, MKLDNN, FP32);\n")
+        file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL_WITH_CUSTOM_TYPE(conv2d, MKLDNN, S8);\n")
+        file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL_WITH_CUSTOM_TYPE(conv2d, MKLDNN, U8);\n")
+        
       else()
         file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL(${TARGET}, MKLDNN);\n")
       endif()

paddle/fluid/framework/block_desc.cc

@@ -13,7 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/framework/block_desc.h"
+
 #include <queue>
+#include <unordered_set>
+#include <utility>
+
 #include "paddle/fluid/framework/operator.h"
 #include "paddle/fluid/framework/program_desc.h"
 
@@ -155,6 +159,16 @@ void BlockDesc::RemoveOp(size_t s, size_t e) {
   ops_.erase(ops_.begin() + s, ops_.begin() + e);
 }
 
+void BlockDesc::RemoveOpInternal(const OpDesc *op_desc) {
+  // TODO(minqiyang): make this faster
+  for (auto it = ops_.begin(); it != ops_.end(); ++it) {
+    if (it->get() == op_desc) {
+      ops_.erase(it);
+      break;
+    }
+  }
+}
+
 std::vector<OpDesc *> BlockDesc::AllOps() const {
   std::vector<OpDesc *> res;
   for (const auto &op : ops_) {
@@ -163,20 +177,6 @@ std::vector<OpDesc *> BlockDesc::AllOps() const {
   return res;
 }
 
-void BlockDesc::Clear() {
-  // clear all ops
-  ops_.clear();
-
-  // clear all vars which are not persistable
-  for (auto it = vars_.begin(); it != vars_.end();) {
-    if (it->second->Persistable()) {
-      ++it;
-    } else {
-      vars_.erase(it++);
-    }
-  }
-}
-
 void BlockDesc::Flush() {
   for (auto &op_desc : ops_) {
     op_desc->Flush();

paddle/fluid/framework/block_desc.h

@@ -93,12 +93,12 @@ class BlockDesc {
    */
   void RemoveOp(size_t s, size_t e);
 
+  void RemoveOpInternal(const OpDesc *op_desc);
+
   void RemoveVar(const std::string &name) { vars_.erase(name); }
 
   std::vector<OpDesc *> AllOps() const;
 
-  void Clear();
-
   size_t OpSize() const { return ops_.size(); }
 
   OpDesc *Op(int idx) const { return ops_.at(idx).get(); }

paddle/fluid/framework/details/fast_threaded_ssa_graph_executor.cc

@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #include "paddle/fluid/framework/details/fast_threaded_ssa_graph_executor.h"
+#include <memory>
 #include <string>
+#include <unordered_map>
 #include <vector>
 #include "paddle/fluid/framework/details/fetch_op_handle.h"
 #include "paddle/fluid/framework/details/multi_devices_helper.h"
@@ -55,7 +57,7 @@ FeedFetchList FastThreadedSSAGraphExecutor::Run(
   std::vector<FetchOpHandle *> fetch_ops;
 
   for (auto &fetch_var_name : fetch_tensors) {
-    for (auto &var_map : graph_->Get<details::GraphVars>("vars")) {
+    for (auto &var_map : graph_->Get<details::GraphVars>(details::kGraphVars)) {
       auto it = var_map.find(fetch_var_name);
       if (it != var_map.end()) {
         fetched_vars[fetch_var_name].push_back(*it->second.rbegin());

paddle/fluid/framework/executor.cc

@@ -20,6 +20,7 @@ limitations under the License. */
 #include "paddle/fluid/framework/lod_tensor_array.h"
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/framework/reader.h"
+#include "paddle/fluid/framework/threadpool.h"
 #include "paddle/fluid/framework/transfer_scope_cache.h"
 #include "paddle/fluid/framework/variable_helper.h"
 #include "paddle/fluid/operators/distributed/distributed.h"

paddle/fluid/framework/ir/CMakeLists.txt

@@ -105,4 +105,5 @@ if (WITH_MKLDNN)
     cc_test(test_conv_bias_mkldnn_fuse_pass SRCS mkldnn/conv_bias_mkldnn_fuse_pass_tester.cc DEPS conv_bias_mkldnn_fuse_pass naive_executor)
     cc_test(test_conv_relu_mkldnn_fuse_pass SRCS mkldnn/conv_relu_mkldnn_fuse_pass_tester.cc DEPS conv_relu_mkldnn_fuse_pass)
     cc_test(test_conv_elementwise_add_mkldnn_fuse_pass SRCS mkldnn/conv_elementwise_add_mkldnn_fuse_pass_tester.cc DEPS conv_elementwise_add_mkldnn_fuse_pass)
+    cc_test(test_mkldnn_placement_pass SRCS mkldnn/mkldnn_placement_pass_tester.cc DEPS mkldnn_placement_pass)
 endif ()

paddle/fluid/framework/ir/mkldnn/conv_elementwise_add_mkldnn_fuse_pass_tester.cc

@@ -44,10 +44,14 @@ struct TestIsReachable {
   using func = std::function<bool(const std::string&, const std::string&)>;
 
   auto operator()(const std::unique_ptr<ir::Graph>& graph) -> func {
-    auto find_node = [](const std::unique_ptr<ir::Graph>& graph,
-                        const std::string& name) -> Node* {
+    auto hash = [](const Node* node) -> std::string {
+      return node->Name() + std::to_string(node->id());
+    };
+
+    auto find_node = [&](const std::unique_ptr<ir::Graph>& graph,
+                         const std::string& name) -> Node* {
       for (auto& node : GraphTraits::DFS(*graph)) {
-        if (name == node.Name()) {
+        if (name == hash(&node)) {
           return &node;
         }
       }
@@ -55,13 +59,17 @@ struct TestIsReachable {
       return nullptr;
     };
 
-    return [&](std::string from, const std::string to) -> bool {
+    // update the from and to strings to hashed equivs in loop from graph traits
+    return [&](std::string from, std::string to) -> bool {
       if (from == to) return true;
 
       std::map<std::string, bool> visited;
 
       for (auto& node : GraphTraits::DFS(*graph)) {
-        visited[node.Name()] = false;
+        auto hashed = hash(&node);
+        if (node.Name() == from) from = hashed;
+        if (node.Name() == to) to = hashed;
+        visited[hashed] = false;
       }
 
       visited[from] = true;
@@ -72,15 +80,15 @@ struct TestIsReachable {
       while (!queue.empty()) {
         auto cur = find_node(graph, queue.front());
         queue.pop_front();
-
         if (cur == nullptr) return false;
 
         for (auto n : cur->outputs) {
-          if (n->Name() == to) return true;
+          auto hashed_name = hash(n);
+          if (hashed_name == to) return true;
 
-          if (!visited[n->Name()]) {
-            visited[n->Name()] = true;
-            queue.push_back(n->Name());
+          if (!visited[hashed_name]) {
+            visited[hashed_name] = true;
+            queue.push_back(hashed_name);
           }
         }
       }
@@ -166,6 +174,28 @@ TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionAsYWithElementwiseAddRelu) {
   RunPassAndAssert(&prog, "a", "relu", 1);
 }
 
+TEST(ConvElementwiseAddMKLDNNFusePass,
+     ConvolutionProjectionAsYWithElementwiseAddRelu) {
+  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e", "f"},
+                               {"bias", "weights", "bias2", "weights2"});
+
+  SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
+  // right branch
+  SetOp(&prog, "conv2d",
+        {{"Input", "b"}, {"Bias", "bias"}, {"Filter", "weights"}},
+        {"Output", "c"});
+
+  // left branch
+  SetOp(&prog, "conv2d",
+        {{"Input", "a"}, {"Bias", "bias2"}, {"Filter", "weights2"}},
+        {"Output", "f"});
+
+  SetOp(&prog, "elementwise_add", {{"X", "f"}, {"Y", "c"}}, {"Out", "d"});
+  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
+
+  RunPassAndAssert(&prog, "a", "relu", 2);
+}
+
 TEST(ConvElementwiseAddMKLDNNFusePass,
      ConvolutionAsYWithElementwiseAddReluNoBias) {
   auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"weights"});

paddle/fluid/framework/ir/mkldnn/mkldnn_placement_pass.cc

@@ -21,7 +21,7 @@ namespace ir {
 
 std::unique_ptr<ir::Graph> MKLDNNPlacementPass::ApplyImpl(
     std::unique_ptr<ir::Graph> graph) const {
-  VLOG(3) << "Aplies MKL-DNN placement strategy.";
+  VLOG(3) << "Applies MKL-DNN placement strategy.";
   const auto& op_types_list =
       Get<std::unordered_set<std::string>>("mkldnn_enabled_op_types");
   for (const Node* n : graph->Nodes()) {

paddle/fluid/framework/ir/mkldnn/mkldnn_placement_pass_tester.cc

+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_placement_pass.h"
+
+#include <gtest/gtest.h>
+#include <boost/logic/tribool.hpp>
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+void SetOp(ProgramDesc* prog, const std::string& type, const std::string& name,
+           const std::vector<std::string>& inputs,
+           const std::vector<std::string>& outputs, boost::tribool use_mkldnn) {
+  auto* op = prog->MutableBlock(0)->AppendOp();
+
+  op->SetType(type);
+
+  if (!boost::indeterminate(use_mkldnn)) op->SetAttr("use_mkldnn", use_mkldnn);
+
+  if (type == "conv2d") {
+    op->SetAttr("name", name);
+    op->SetInput("Input", {inputs[0]});
+    op->SetInput("Filter", {inputs[1]});
+    op->SetInput("Bias", {inputs[2]});
+  } else if (type == "relu") {
+    op->SetInput("X", inputs);
+  } else if (type == "concat") {
+    op->SetAttr("axis", 1);
+    op->SetInput("X", {inputs[0], inputs[1]});
+  } else if (type == "pool2d") {
+    op->SetInput("X", {inputs[0]});
+  } else {
+    FAIL() << "Unexpected operator type.";
+  }
+  op->SetOutput("Out", {outputs[0]});
+}
+
+// operator                      use_mkldnn
+// ---------------------------------------
+// (a,b)->concat->c              none
+// (c,weights,bias)->conv->f     none
+// f->relu->g                    false
+// g->pool->h                    false
+// (h,weights2,bias2)->conv->k   true
+// k->relu->l                    true
+ProgramDesc BuildProgramDesc() {
+  ProgramDesc prog;
+
+  for (auto& v :
+       std::vector<std::string>({"a", "b", "c", "weights", "bias", "f", "g",
+                                 "h", "weights2", "bias2", "k", "l"})) {
+    auto* var = prog.MutableBlock(0)->Var(v);
+    var->SetType(proto::VarType::SELECTED_ROWS);
+    if (v == "weights" || v == "bias") {
+      var->SetPersistable(true);
+    }
+  }
+
+  SetOp(&prog, "concat", "concat1", std::vector<std::string>({"a", "b"}),
+        std::vector<std::string>({"c"}), boost::indeterminate);
+  SetOp(&prog, "conv2d", "conv1",
+        std::vector<std::string>({"c", "weights", "bias"}),
+        std::vector<std::string>({"f"}), boost::indeterminate);
+  SetOp(&prog, "relu", "relu1", std::vector<std::string>({"f"}),
+        std::vector<std::string>({"g"}), false);
+  SetOp(&prog, "pool2d", "pool1", std::vector<std::string>({"g"}),
+        std::vector<std::string>({"h"}), false);
+  SetOp(&prog, "conv2d", "conv2",
+        std::vector<std::string>({"h", "weights2", "bias2"}),
+        std::vector<std::string>({"k"}), true);
+  SetOp(&prog, "relu", "relu2", std::vector<std::string>({"k"}),
+        std::vector<std::string>({"l"}), true);
+
+  return prog;
+}
+
+void MainTest(std::initializer_list<std::string> mkldnn_enabled_op_types,
+              unsigned expected_use_mkldnn_true_count) {
+  auto prog = BuildProgramDesc();
+
+  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+
+  auto pass = PassRegistry::Instance().Get("mkldnn_placement_pass");
+  pass->Set("mkldnn_enabled_op_types",
+            new std::unordered_set<std::string>(mkldnn_enabled_op_types));
+
+  graph = pass->Apply(std::move(graph));
+
+  unsigned use_mkldnn_true_count = 0;
+
+  for (auto* node : graph->Nodes()) {
+    if (node->IsOp()) {
+      auto* op = node->Op();
+      if (op->HasAttr("use_mkldnn") &&
+          boost::get<bool>(op->GetAttr("use_mkldnn"))) {
+        ++use_mkldnn_true_count;
+      }
+    }
+  }
+
+  EXPECT_EQ(use_mkldnn_true_count, expected_use_mkldnn_true_count);
+}
+
+TEST(MKLDNNPlacementPass, enable_conv_relu) {
+  // 1 conv (1 conv is always true) + 2 relu (1 relu is always true) + 0 pool
+  MainTest({"conv2d", "relu"}, 3);
+}
+
+TEST(MKLDNNPlacementPass, enable_relu_pool) {
+  // 1 conv (1 conv is always true) + 2 relu (1 relu is always true) + 1 pool
+  MainTest({"relu", "pool2d"}, 4);
+}
+
+TEST(MKLDNNPlacementPass, enable_all) {
+  // 1 conv (1 conv is always true) + 2 relu (1 relu is always true) + 1 pool
+  MainTest({}, 4);
+}
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+USE_PASS(mkldnn_placement_pass);

paddle/fluid/framework/op_registry.h

@@ -290,7 +290,7 @@ struct OpKernelRegistrarFunctorEx<PlaceType, false, I,
       "USE_OP_DEVICE_KERNEL must be in global namespace");                 \
   extern int                                                               \
       TouchOpKernelRegistrar_##op_type##_##LIBRARY_TYPE##_##customized_name(); \
-  UNUSED static int use_op_kernel_##op_type##_##LIBRARY_TYPE##_##DEFAULT_TYPE##_ = /* NOLINT */ \
+  UNUSED static int use_op_kernel_##op_type##_##LIBRARY_TYPE##_##customized_name##_ = /* NOLINT */ \
       TouchOpKernelRegistrar_##op_type##_##LIBRARY_TYPE##_##customized_name()
 
 #define USE_OP_DEVICE_KERNEL(op_type, LIBRARY_TYPE) \

paddle/fluid/framework/operator.cc

@@ -882,7 +882,8 @@ class RuntimeInferShapeContext : public InferShapeContext {
   const RuntimeContext& ctx_;
 };
 
-static void CheckTensorNANOrInf(const std::string& name,
+static void CheckTensorNANOrInf(const std::string& op_type,
+                                const std::string& name,
                                 const framework::Tensor& tensor) {
   if (tensor.memory_size() == 0) {
     return;
@@ -892,9 +893,9 @@ static void CheckTensorNANOrInf(const std::string& name,
     return;
   }
   PADDLE_ENFORCE(!framework::TensorContainsInf(tensor),
-                 "Tensor %s contains Inf", name);
+                 "Operator %s output Tensor %s contains Inf", op_type, name);
   PADDLE_ENFORCE(!framework::TensorContainsNAN(tensor),
-                 "Tensor %s contains NAN", name);
+                 "Operator %s output Tensor %s contains NAN", op_type, name);
 }
 
 void OperatorWithKernel::RuntimeInferShape(const Scope& scope,
@@ -988,9 +989,10 @@ void OperatorWithKernel::RunImpl(const Scope& scope,
       auto* var = exec_scope.FindVar(vname);
       if (var == nullptr) continue;
       if (var->IsType<framework::LoDTensor>()) {
-        CheckTensorNANOrInf(vname, var->Get<framework::LoDTensor>());
+        CheckTensorNANOrInf(type_, vname, var->Get<framework::LoDTensor>());
       } else if (var->IsType<framework::SelectedRows>()) {
-        CheckTensorNANOrInf(vname, var->Get<framework::SelectedRows>().value());
+        CheckTensorNANOrInf(type_, vname,
+                            var->Get<framework::SelectedRows>().value());
       }
     }
   }

paddle/fluid/framework/python_headers.h

@@ -24,3 +24,11 @@ limitations under the License. */
 
 #pragma pop_macro("_XOPEN_SOURCE")
 #pragma pop_macro("_POSIX_C_SOURCE")
+
+#if !defined(PYBIND11_HIDDEN)
+#ifdef _WIN32
+#define PYBIND11_HIDDEN __declspec(dllexport)
+#else
+#define PYBIND11_HIDDEN __attribute__((visibility("hidden")))
+#endif
+#endif

paddle/fluid/framework/tensor_util.cc

@@ -14,6 +14,8 @@
 #include "paddle/fluid/framework/tensor_util.h"
 #include <algorithm>
 #include <limits>
+#include <memory>
+#include <utility>
 #include <vector>
 #include "paddle/fluid/framework/data_type.h"
 

paddle/fluid/imperative/layer.cc

@@ -18,6 +18,7 @@
 #include <limits>
 #include <map>
 #include <random>
+#include <unordered_set>
 #include <utility>
 
 #include "paddle/fluid/framework/lod_tensor.h"
@@ -139,6 +140,8 @@ class Autograd {
           }
         }
       }
+
+      ready_op->InvokeBackwardHooks();
     }
   }
 
@@ -156,8 +159,10 @@ class Autograd {
       for (auto it : candidate->pre_ops_) {
         for (OpBase* pre_op : it.second) {
           if (!pre_op) continue;
-          VLOG(5) << "op dep " << candidate->op_desc_->Type() << " <---- "
-                  << it.first << " <---- " << pre_op->op_desc_->Type();
+          VLOG(5) << "op dep " << candidate->op_desc_->Type() << " trace id "
+                  << candidate->trace_id_ << " <---- " << it.first << " <---- "
+                  << pre_op->op_desc_->Type() << " trace id "
+                  << pre_op->trace_id_;
           if (visited.find(pre_op) == visited.end()) {
             visited.insert(pre_op);
             queue.push_back(pre_op);
@@ -211,6 +216,7 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
     return {};
   }
 
+  VLOG(3) << "apply op grad: " << op_desc_->Type();
   std::vector<framework::VariableValueMap> grad_outputs;
   if (backward_id_ > 0) {
     VLOG(3) << "py_layer_grad";
@@ -272,6 +278,22 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
   return input_vars_;
 }
 
+void OpBase::InvokeBackwardHooks() {
+  VLOG(3) << "call backward hooks, hooks num: " << backward_hooks_.size();
+
+  // call backward hooks
+  for (py::object& callable : backward_hooks_) {
+    callable(this);
+  }
+}
+
+void OpBase::RegisterBackwardHooks(const py::object& callable) {
+  VLOG(3) << "Register backward hooks " << trace_id_;
+
+  // TODO(minqiyang): check the callable format
+  backward_hooks_.push_back(callable);
+}
+
 void VarBase::RunBackward() {
   if (!pre_op_) return;
 

paddle/fluid/imperative/layer.h

@@ -123,22 +123,32 @@ class VarBase {
 
  private:
   VarBase(framework::Variable* var, VarBase* grad, bool stop_gradient)
-      : var_desc_(nullptr),
+      : name_(),
+        var_desc_(nullptr),
         var_(var),
         grads_(grad),
+        block_(nullptr),
+        persistable_(false),
         stop_gradient_(stop_gradient),
         pre_op_(nullptr),
+        pre_op_out_name_(),
         pre_op_out_idx_(-1) {}
 
  public:
   virtual ~VarBase() {
+    // TODO(minqiyang): remove var desc from block desc
     if (var_) {
       delete var_;
+      var_ = nullptr;
     }
 
     if (grads_) {
       delete grads_;
+      grads_ = nullptr;
     }
+
+    pre_op_ = nullptr;
+    pre_op_out_idx_ = -1;
   }
 
   inline OpBase* PreOp() const { return pre_op_; }
@@ -151,6 +161,14 @@ class VarBase {
 
   void RunBackward();
 
+  inline void ResetPreOp(OpBase* op) {
+    if (op == pre_op_) {
+      // clear pre_op info when op equals to var's pre_op
+      pre_op_ = nullptr;
+      pre_op_out_idx_ = -1;
+    }
+  }
+
   void TrackPreOp(OpBase* pre_op, const std::string& pre_op_out_name,
                   int pre_op_out_idx, bool pre_op_stop_gradient) {
     pre_op_ = pre_op;
@@ -184,11 +202,15 @@ class VarBase {
     return string::Sprintf("%s@IGrad", var_desc_->Name());
   }
 
+  std::string name_;
   framework::VarDesc* var_desc_;
 
   framework::Variable* var_;
   VarBase* grads_;
 
+  framework::BlockDesc* block_;
+  bool persistable_;
+
  private:
   bool stop_gradient_;
   OpBase* pre_op_;
@@ -199,15 +221,27 @@ class VarBase {
 /* The wrapper for OpDesc which holds a OpDesc and a OpDesc of its
  * gradient. This object should be managed totally by Python intepreter.
  */
-class OpBase {
+class PYBIND11_HIDDEN OpBase {
  public:
   OpBase()
       : op_desc_(nullptr),
         forward_id_(-1),
         backward_id_(-1),
-        place_(platform::CPUPlace()) {}
+        trace_id_(-1),
+        place_(platform::CPUPlace()),
+        backward_hooks_() {}
 
   virtual ~OpBase() {
+    // TODO(minqiyang): remove op_desc from block_desc in tracer
+    //
+    // reset all output vars' pre op
+    for (auto iter : output_vars_) {
+      for (VarBase* var : iter.second) {
+        var->ResetPreOp(this);
+      }
+    }
+
+    // release resource
     for (framework::OpDesc* desc : grad_op_descs_) {
       delete desc;
     }
@@ -215,6 +249,10 @@ class OpBase {
 
   std::map<std::string, std::vector<VarBase*>> ApplyGrad();
 
+  void RegisterBackwardHooks(const py::object& callable);
+
+  void InvokeBackwardHooks();
+
   // One of `op_desc_` or `forward_id_` is set, not both.
   // For pure python PyLayer, use `forward_id_`, otherwise, use op_desc_.
   framework::OpDesc* op_desc_;
@@ -225,6 +263,7 @@ class OpBase {
   // Note: each fwd op corresponds to a vector of bwd ops.
   std::vector<framework::OpDesc*> grad_op_descs_;
   int backward_id_;
+  int trace_id_;
 
   platform::Place place_;
 
@@ -239,6 +278,8 @@ class OpBase {
   std::vector<framework::VariableValueMap> grad_output_vars_;
 
   framework::BlockDesc* block_;
+
+  std::vector<py::object> backward_hooks_;
 };
 
 class Layer {

paddle/fluid/imperative/tracer.cc

@@ -14,15 +14,32 @@
 
 #include "paddle/fluid/imperative/tracer.h"
 
+#include <memory>
 #include <set>
+#include <unordered_map>
+#include <unordered_set>
 
 #include "paddle/fluid/operators/math/math_function.h"
 #include "paddle/fluid/platform/device_context.h"
 #include "paddle/fluid/platform/enforce.h"
 
+#ifdef WITH_GPERFTOOLS
+#include "gperftools/profiler.h"
+#endif
+
+DEFINE_string(
+    tracer_profile_fname, "",
+    "Profiler filename for imperative tracer, which generated by gperftools."
+    "Only valid when compiled `WITH_PROFILER=ON`. Empty if disable.");
+
 namespace paddle {
 namespace imperative {
 
+static std::once_flag gTracerProfileOnce;
+#ifdef WITH_GPERFTOOLS
+static bool gTracerProfilerStarted = false;
+#endif
+
 void CreateGradOp(const framework::OpDesc& op_desc,
                   const std::unordered_set<std::string>& no_grad_set,
                   const std::vector<framework::BlockDesc*>& grad_sub_block,
@@ -68,15 +85,36 @@ platform::Place GetExpectedPlace(platform::Place place, VarBasePtrMap inputs) {
   return result;
 }
 
+Tracer::Tracer(framework::BlockDesc* root_block) : root_block_(root_block) {
+  if (!FLAGS_tracer_profile_fname.empty()) {
+    std::call_once(gTracerProfileOnce, [] {
+#ifdef WITH_GPERFTOOLS
+      ProfilerStart(FLAGS_tracer_profile_fname.c_str());
+      gTracerProfilerStarted = true;
+#else
+      LOG(WARNING) << "Paddle is not compiled with gperftools. "
+                      "FLAGS_tracer_profile_fname will be ignored";
+#endif
+    });
+  }
+}
+
 std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
                                     const VarBasePtrMap& outputs,
                                     framework::BlockDesc* block,
                                     const platform::Place expected_place,
                                     const bool stop_gradient) {
+#ifdef WITH_GPERFTOOLS
+  if (gTracerProfilerStarted) {
+    ProfilerFlush();
+  }
+#endif
+
   std::map<std::string, VarBase*> vars;
 
   framework::OpDesc* op_desc = op->op_desc_;
-  VLOG(3) << "tracer tracing " << op_desc->Type();
+  VLOG(3) << "tracer tracing " << op_desc->Type() << " trace id "
+          << op->trace_id_;
   op_desc->InferShape(*block);
   op_desc->InferVarType(block);
 
@@ -99,11 +137,13 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
       if (inp->PreOp() && !inp->IsStopGradient()) {
         op->pre_ops_[it.first].push_back(inp->PreOp());
         op->pre_ops_out_idx_[it.first].push_back(inp->PreOpOutIdx());
+        VLOG(3) << "add pre op " << inp->PreOp()->op_desc_->Type();
       } else {
         op->pre_ops_[it.first].push_back(nullptr);
       }
       VLOG(3) << "input vname " << inp->var_desc_->Name() << " "
-              << inp->var_->IsInitialized();
+              << inp->var_->IsInitialized() << " stop_gradient "
+              << inp->IsStopGradient();
     }
   }
 
@@ -155,6 +195,7 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
 
     op->grad_input_vars_.resize(op->grad_op_descs_.size());
     op->grad_output_vars_.resize(op->grad_op_descs_.size());
+
     for (size_t i = 0; i < op->grad_op_descs_.size(); ++i) {
       framework::OpDesc* grad_op_desc = op->grad_op_descs_[i];
       for (auto it : grad_op_desc->Inputs()) {
@@ -167,7 +208,6 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
             PADDLE_ENFORCE(fwd_var_it != vars.end());
             // Forward inputs or outputs.
             grad_in_vars.push_back(fwd_var_it->second->var_);
-            vars_saved_for_backward.insert(it.first);
           } else {
             VarBase* var = vars[var_it->second];
             if (!var->grads_->var_->IsInitialized()) {
@@ -177,6 +217,8 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
             // Douts.
             grad_in_vars.push_back(var->grads_->var_);
           }
+
+          vars_saved_for_backward.insert(it.first);
         }
       }
 

paddle/fluid/imperative/tracer.h

@@ -40,7 +40,7 @@ platform::Place GetExpectedPlace(platform::Place place, VarBasePtrMap inputs);
 
 class Tracer {
  public:
-  explicit Tracer(framework::BlockDesc* root_block) : root_block_(root_block) {}
+  explicit Tracer(framework::BlockDesc* root_block);
 
   virtual ~Tracer() {}
 

paddle/fluid/inference/CMakeLists.txt

@@ -16,6 +16,7 @@ add_subdirectory(utils)
 if (TENSORRT_FOUND)
   add_subdirectory(tensorrt)
 endif()
+# add_subdirectory(anakin)
 
 get_property(fluid_modules GLOBAL PROPERTY FLUID_MODULES)
 get_property(cuda_modules GLOBAL PROPERTY CUDA_MODULES)

paddle/fluid/inference/anakin/CMakeLists.txt

+cc_library(anakin_engine SRCS engine.cc)
+target_link_libraries(anakin_engine anakin anakin_saber_common)
+cc_test(test_anakin_engine SRCS test_anakin_engine.cc DEPS anakin_engine)
+add_subdirectory(convert)

paddle/fluid/inference/anakin/convert/CMakeLists.txt

+cc_library(anakin_op_converter SRCS fc.cc registrar.cc DEPS anakin_engine framework_proto scope)
+cc_test(test_anakin_fc SRCS test_fc_op.cc DEPS anakin_op_converter mul_op)

paddle/fluid/inference/anakin/convert/fc.cc

+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/inference/anakin/convert/fc.h"
+#include <algorithm>
+
+using anakin::graph::GraphGlobalMem;
+using anakin::AK_FLOAT;
+using anakin::Precision;
+using anakin::saber::NV;
+using anakin::saber::X86;
+using anakin::saber::Shape;
+using anakin::PBlock;
+using anakin::PTuple;
+
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+void FcOpConverter::operator()(const framework::proto::OpDesc &op,
+                               const framework::Scope &scope, bool test_mode) {
+  framework::OpDesc op_desc(op, nullptr);
+  PADDLE_ENFORCE_EQ(op_desc.Input("X").size(), 1);
+  PADDLE_ENFORCE_EQ(op_desc.Input("Y").size(), 1);
+  PADDLE_ENFORCE_EQ(op_desc.Output("Out").size(), 1);
+
+  auto x_name = op_desc.Input("X").front();
+  auto op_name = op_desc.Type() + ":" + op_desc.Output("Out").front();
+  auto *y_v = scope.FindVar(op_desc.Input("Y").front());
+  PADDLE_ENFORCE_NOT_NULL(y_v);
+  auto *y_t = y_v->GetMutable<framework::LoDTensor>();
+
+  auto input_name = op_desc.Input("X").front();
+  auto output_name = op_desc.Output("Out").front();
+
+  auto weight_shape = framework::vectorize2int(y_t->dims());
+  engine_->AddOp(op_name, "Dense", {input_name}, {output_name});
+  engine_->AddOpAttr(op_name, "bias_term", false);
+  engine_->AddOpAttr(op_name, "axis", 1);
+  int out_dim = weight_shape[1];
+  engine_->AddOpAttr(op_name, "out_dim", out_dim);
+
+  weight_shape.push_back(1);
+  weight_shape.push_back(1);
+  Shape anakin_shape(weight_shape);
+
+  framework::LoDTensor weight_tensor;
+  weight_tensor.Resize(y_t->dims());
+  TensorCopySync((*y_t), platform::CPUPlace(), &weight_tensor);
+
+  auto *weight1 =
+      GraphGlobalMem<NV>::Global().template new_block<AK_FLOAT>(anakin_shape);
+  float *cpu_data = static_cast<float *>(weight1->h_tensor().mutable_data());
+  std::copy_n(weight_tensor.data<float>(), weight_tensor.numel(), cpu_data);
+  weight1->d_tensor().set_shape(anakin_shape);
+  weight1->d_tensor().copy_from(weight1->h_tensor());
+  engine_->AddOpAttr(op_name, "weight_1", *weight1);
+}
+
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/anakin/convert/fc.h

+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/fluid/inference/anakin/convert/op_converter.h"
+
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+class FcOpConverter : public AnakinOpConverter {
+ public:
+  FcOpConverter() = default;
+
+  virtual void operator()(const framework::proto::OpDesc &op,
+                          const framework::Scope &scope,
+                          bool test_mode) override;
+  virtual ~FcOpConverter() {}
+
+ private:
+};
+
+static Registrar<FcOpConverter> register_fc_op_converter("fc");
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/anakin/convert/op_converter.h

+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include "framework/core/types.h"
+#include "paddle/fluid/framework/block_desc.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/inference/anakin/convert/registrar.h"
+#include "paddle/fluid/inference/anakin/engine.h"
+#include "paddle/fluid/inference/utils/singleton.h"
+#include "saber/saber_types.h"
+
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+using AnakinNvEngine =
+    AnakinEngine<::anakin::saber::NV, ::anakin::Precision::FP32>;
+
+class AnakinOpConverter {
+ public:
+  AnakinOpConverter() = default;
+
+  virtual void operator()(const framework::proto::OpDesc &op,
+                          const framework::Scope &scope, bool test_mode) {}
+  void ConvertOp(const framework::proto::OpDesc &op,
+                 const std::unordered_set<std::string> &parameters,
+                 const framework::Scope &scope, AnakinNvEngine *engine,
+                 bool test_mode = false) {
+    framework::OpDesc op_desc(op, nullptr);
+    std::string op_type = op_desc.Type();
+    std::shared_ptr<AnakinOpConverter> it{nullptr};
+
+    if (op_type == "mul") {
+      PADDLE_ENFORCE_EQ(op_desc.Input("Y").size(), 1UL);
+      std::string Y = op_desc.Input("Y")[0];
+      std::cout << Y << parameters.count(Y) << std::endl;
+      if (parameters.count(Y)) {
+        it = OpRegister::instance()->Get("fc");
+      }
+    }
+
+    if (!it) {
+      it = OpRegister::instance()->Get(op_type);
+    }
+    PADDLE_ENFORCE_NOT_NULL(it, "no OpConverter for optype [%s]", op_type);
+    it->SetEngine(engine);
+    (*it)(op, scope, test_mode);
+  }
+
+  void ConvertBlock(const framework::proto::BlockDesc &block,
+                    const std::unordered_set<std::string> &parameters,
+                    const framework::Scope &scope, AnakinNvEngine *engine) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    for (auto i = 0; i < block.ops_size(); i++) {
+      auto &op = block.ops(i);
+      ConvertOp(op, parameters, scope, engine);
+    }
+  }
+  void SetEngine(AnakinNvEngine *engine) { engine_ = engine; }
+  virtual ~AnakinOpConverter() {}
+
+ protected:
+  bool test_mode_;
+  AnakinNvEngine *engine_{nullptr};
+
+ private:
+  std::unordered_map<std::string, AnakinOpConverter *> converters_;
+  framework::Scope *scope_{nullptr};
+  std::mutex mutex_;
+};
+
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle
+
+#define REGISTER_ANAKIN_OP_CONVERTER(op_type__, Converter__)                \
+  struct anakin_##op_type__##_converter                                     \
+      : public ::paddle::framework::Registrar {                             \
+    anakin_##op_type__##_converter() {                                      \
+      ::paddle::inference::                                                 \
+          Registry<paddle::inference::anakin::AnakinOpConverter>::Register< \
+              ::paddle::inference::anakin::Converter__>(#op_type__);        \
+    }                                                                       \
+  };                                                                        \
+  anakin_##op_type__##_converter anakin_##op_type__##_converter__;          \
+  int TouchConverterRegister_anakin_##op_type__() {                         \
+    anakin_##op_type__##_converter__.Touch();                               \
+    return 0;                                                               \
+  }
+
+#define USE_ANAKIN_CONVERTER(op_type__)                                    \
+  extern int TouchConverterRegister_anakin_##op_type__();                  \
+  static int use_op_converter_anakin_##op_type__ __attribute__((unused)) = \
+      TouchConverterRegister_anakin_##op_type__();

paddle/fluid/inference/anakin/convert/registrar.cc

+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/inference/anakin/convert/registrar.h"
+
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+std::shared_ptr<AnakinOpConverter> OpRegister::Get(const std::string &name) {
+  auto it = registry_.find(name);
+  if (it == registry_.end()) return nullptr;
+  return it->second();
+}
+
+OpRegister *OpRegister::instance() {
+  static OpRegister factory;
+  return &factory;
+}
+
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/anakin/convert/registrar.h

+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <functional>
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+class AnakinOpConverter;
+
+class OpRegister {
+ public:
+  OpRegister() = default;
+  std::shared_ptr<AnakinOpConverter> Get(const std::string &name);
+  static OpRegister *instance();
+  void OpRegisterFn(const std::string &name,
+                    std::function<std::shared_ptr<AnakinOpConverter>()> fn) {
+    registry_[name] = fn;
+  }
+
+ private:
+  using RegisterFnType = std::function<std::shared_ptr<AnakinOpConverter>()>;
+  std::map<std::string, std::function<std::shared_ptr<AnakinOpConverter>()>>
+      registry_;
+};
+
+template <typename T, typename... Args>
+class Registrar {
+ public:
+  Registrar(const std::string &name, Args... args) {
+    std::shared_ptr<AnakinOpConverter> converter =
+        std::make_shared<T>(std::move(args)...);
+    OpRegister::instance()->OpRegisterFn(name,
+                                         [converter]() { return converter; });
+  }
+};
+
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/anakin/convert/test_fc_op.cc

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <gtest/gtest.h>
+#include "paddle/fluid/inference/anakin/convert/fc.h"
+#include "paddle/fluid/inference/anakin/convert/op_converter.h"
+#include "paddle/fluid/inference/anakin/convert/ut_helper.h"
+
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+TEST(fc_op, test) {
+  auto fc_converter = OpRegister::instance()->Get("fc");
+  ASSERT_TRUE(fc_converter != nullptr);
+  // Registrar<FcOpConverter> register_fc("fc");
+  // auto fc = std::make_shared<FcOpConverter>();
+
+  std::unordered_set<std::string> parameters({"mul_y"});
+  framework::Scope scope;
+  AnakinConvertValidation validator(parameters, scope);
+  validator.DeclInputVar("mul_x", {1, 1, 1, 1});
+  validator.DeclParamVar("mul_y", {1, 2});
+  validator.DeclOutputVar("mul_out", {1, 1, 1, 2});
+
+  // Prepare Op description
+  framework::OpDesc desc;
+  desc.SetType("mul");
+  desc.SetInput("X", {"mul_x"});
+  desc.SetInput("Y", {"mul_y"});
+  desc.SetOutput("Out", {"mul_out"});
+  int num_flatten_dims = 3;
+  desc.SetAttr("x_num_col_dims", num_flatten_dims);
+  validator.SetOp(*desc.Proto());
+
+  validator.Execute(10);
+}
+
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle
+
+USE_OP(mul);

paddle/fluid/inference/anakin/convert/ut_helper.h

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include <map>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "paddle/fluid/framework/lod_tensor.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/anakin/engine.h"
+#include "paddle/fluid/inference/analysis/helper.h"
+#include "paddle/fluid/inference/utils/singleton.h"
+#include "paddle/fluid/platform/enforce.h"
+
+using anakin::graph::GraphGlobalMem;
+using anakin::AK_FLOAT;
+using anakin::Precision;
+using anakin::saber::NV;
+using anakin::saber::X86;
+using anakin::saber::Shape;
+using anakin::PBlock;
+using anakin::PTuple;
+
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+/*
+ * Get a random float value between [low, high]
+ */
+float random(float low, float high) {
+  static std::random_device rd;
+  static std::mt19937 mt(rd());
+  std::uniform_real_distribution<double> dist(low, high);
+  return dist(mt);
+}
+
+void RandomizeTensor(framework::LoDTensor* tensor, const platform::Place& place,
+                     const platform::DeviceContext& ctx) {
+  auto dims = tensor->dims();
+  size_t num_elements = analysis::AccuDims(dims, dims.size());
+  PADDLE_ENFORCE_GT(num_elements, 0);
+
+  platform::CPUPlace cpu_place;
+  framework::LoDTensor temp_tensor;
+  temp_tensor.Resize(dims);
+  auto* temp_data = temp_tensor.mutable_data<float>(cpu_place);
+
+  for (size_t i = 0; i < num_elements; i++) {
+    *(temp_data + i) = random(0., 1.);
+  }
+
+  TensorCopySync(temp_tensor, place, tensor);
+}
+
+/*
+ * Help to validate the correctness between Fluid Op and the corresponding
+ * anakin
+ * layer.
+ */
+class AnakinConvertValidation {
+  using AnakinNvEngineT = AnakinEngine<NV, Precision::FP32>;
+
+ public:
+  AnakinConvertValidation() = delete;
+
+  AnakinConvertValidation(const std::unordered_set<std::string>& parameters,
+                          const framework::Scope& scope)
+      : parameters_(parameters), scope_(scope), place_(0) {
+    PADDLE_ENFORCE_EQ(cudaStreamCreate(&stream_), 0);
+    engine_.reset(new AnakinEngine<NV, Precision::FP32>(true));
+  }
+
+  // Declare a Variable as input with random initialization.
+  void DeclInputVar(const std::string& name,
+                    const std::vector<int> tensor_dims) {
+    DeclVar(name, tensor_dims);
+    // should decalre anakin input here.
+  }
+
+  void DeclParamVar(const std::string& name, const std::vector<int> dim_vec) {
+    DeclVar(name, dim_vec);
+  }
+
+  void DeclOutputVar(const std::string& name, const std::vector<int> dim_vec) {
+    DeclVar(name, dim_vec);
+    // should declare anakin output here.
+  }
+
+  void DeclVar(const std::string& name, const std::vector<int> dim_vec) {
+    platform::CUDADeviceContext ctx(place_);
+    auto* x = scope_.Var(name);
+    auto* x_tensor = x->GetMutable<framework::LoDTensor>();
+    x_tensor->Resize(framework::make_ddim(dim_vec));
+    RandomizeTensor(x_tensor, place_, ctx);
+  }
+
+  void SetOp(const framework::proto::OpDesc& desc) {
+    op_ = framework::OpRegistry::CreateOp(desc);
+    op_desc_.reset(new framework::OpDesc(desc, nullptr));
+    // should init anakin engine here.
+
+    Singleton<AnakinOpConverter>::Global().ConvertOp(
+        desc, parameters_, scope_, engine_.get(), true /*test_mode*/);
+    engine_->Freeze();
+    for (const auto& input : op_desc_->InputArgumentNames()) {
+      if (parameters_.count(input)) continue;
+      auto& t = inference::analysis::GetFromScope<framework::LoDTensor>(scope_,
+                                                                        input);
+      auto t_shape = framework::vectorize2int(t.dims());
+      engine_->SetInputShape(input, t_shape);
+    }
+    engine_->Optimize();
+    engine_->InitGraph();
+  }
+
+  // We use the set 'neglected_output' here, because some Ops like batch norm,
+  // the outputs specified in the op des are only used during training,
+  // so we should neglect those output during inference.
+  void Execute(int batch_size,
+               std::unordered_set<std::string> neglected_output = {}) {
+    // Execute Fluid Op
+    platform::CUDADeviceContext ctx(place_);
+    op_->Run(scope_, place_);
+
+    // std::vector<framework::LoDTensor> input_vector;
+    // std::vector<framework::LoDTensor> output_vector;
+    std::map<std::string, framework::LoDTensor*> inputs;
+    for (const auto& input : op_desc_->InputArgumentNames()) {
+      if (parameters_.count(input)) continue;
+      auto* var = scope_.FindVar(input);
+      auto tensor = var->GetMutable<framework::LoDTensor>();
+      inputs.insert({input, tensor});
+    }
+
+    std::map<std::string, framework::LoDTensor*> outputs;
+    std::vector<std::vector<float>> fluid_outputs;
+    for (const auto& output : op_desc_->OutputArgumentNames()) {
+      if (neglected_output.count(output)) continue;
+      std::vector<float> fluid_out;
+      auto* var = scope_.FindVar(output);
+      auto tensor = var->GetMutable<framework::LoDTensor>();
+      framework::TensorToVector(*tensor, ctx, &fluid_out);
+      fluid_outputs.push_back(fluid_out);
+
+      // size_t fluid_out_size = fluid_out.size();
+      /*for (size_t i = 0; i < fluid_out_size; i++) {
+        std::cout << fluid_out[i] << std::endl;
+      }*/
+      outputs.insert({output, tensor});
+    }
+
+    engine_->Execute(inputs, outputs);
+    int i_output = 0;
+    for (const auto& output : op_desc_->OutputArgumentNames()) {
+      if (neglected_output.count(output)) continue;
+      std::vector<float> anakin_out;
+      auto* var = scope_.FindVar(output);
+      auto tensor = var->GetMutable<framework::LoDTensor>();
+      framework::TensorToVector(*tensor, ctx, &anakin_out);
+
+      size_t anakin_out_size = anakin_out.size();
+      auto fluid_out = fluid_outputs[i_output++];
+      for (size_t i = 0; i < anakin_out_size; i++) {
+        LOG(INFO) << "Output[" << i << "]: anakin[" << anakin_out[i] << "], "
+                  << "fluid[" << fluid_out[i] << "]";
+      }
+    }
+  }
+
+  framework::Scope& scope() { return scope_; }
+
+ private:
+  std::unique_ptr<AnakinNvEngineT> engine_{nullptr};
+  cudaStream_t stream_;
+  std::unique_ptr<framework::OperatorBase> op_;
+  std::unique_ptr<framework::OpDesc> op_desc_;
+  const std::unordered_set<std::string>& parameters_;
+  framework::Scope& scope_;
+  platform::CUDAPlace place_;
+};
+
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/anakin/engine.cc

+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/inference/anakin/engine.h"
+#include <algorithm>
+#include <cstring>
+#include <map>
+#include <utility>
+#include "paddle/fluid/framework/ddim.h"
+
+using anakin::Precision;
+using anakin::OpRunType;
+using paddle::framework::LoDTensor;
+template <typename T, Precision P, OpRunType O>
+using AnakinNetT = anakin::Net<T, P, O>;
+
+template <typename T, Precision P>
+using AnakinGraphT = anakin::graph::Graph<T, P>;
+
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+template <typename TargetT, Precision PrecisionType, OpRunType RunType>
+AnakinEngine<TargetT, PrecisionType, RunType>::AnakinEngine(bool need_summary)
+    : graph_(new AnakinGraphT<TargetT, PrecisionType>()),
+      net_(new AnakinNetT<TargetT, PrecisionType, RunType>(need_summary)) {}
+
+template <typename TargetT, Precision PrecisionType, OpRunType RunType>
+AnakinEngine<TargetT, PrecisionType, RunType>::~AnakinEngine() {}
+
+template <typename TargetT, Precision PrecisionType, OpRunType RunType>
+void AnakinEngine<TargetT, PrecisionType, RunType>::SetInputShape(
+    const std::string &name, std::vector<int> shape) {
+  graph_->AddOpAttr<::anakin::PTuple<int>>(name, "input_shape",
+                                           std::move(shape));
+}
+
+template <typename TargetT, Precision PrecisionType, OpRunType RunType>
+void AnakinEngine<TargetT, PrecisionType, RunType>::InitGraph() {
+  net_->init(*graph_);
+}
+
+template <typename TargetT, Precision PrecisionType, OpRunType RunType>
+void AnakinEngine<TargetT, PrecisionType, RunType>::AddOp(
+    const std::string &name, const std::string &type,
+    const std::vector<std::string> &inputs,
+    const std::vector<std::string> &outputs) {
+  PADDLE_ENFORCE(graph_->AddOp(name, type, inputs, outputs), "Add operation.");
+}
+
+template <typename TargetT, Precision PrecisionType, OpRunType RunType>
+void AnakinEngine<TargetT, PrecisionType, RunType>::Execute(
+    const std::map<std::string, framework::LoDTensor *> &inputs,
+    const std::map<std::string, framework::LoDTensor *> &outputs) {
+  for (const auto &input : inputs) {
+    auto *tensor = input.second;
+    auto *data = tensor->data<float>();
+    auto shape = framework::vectorize2int(tensor->dims());
+    ::anakin::saber::Shape anakin_shape(shape);
+    auto *anakin_input = net_->get_in(input.first);
+    ::anakin::saber::Tensor<TargetT> tmp_anakin_tensor(data, TargetT(), 0,
+                                                       anakin_shape);
+    anakin_input->share_from(tmp_anakin_tensor);
+  }
+
+  for (const auto &output : outputs) {
+    auto *tensor = output.second;
+    auto *data = tensor->data<float>();
+    auto shape = framework::vectorize2int(tensor->dims());
+    ::anakin::saber::Shape anakin_shape(shape);
+    auto *anakin_output = net_->get_out(output.first);
+    ::anakin::saber::Tensor<TargetT> tmp_anakin_tensor(data, TargetT(), 0,
+                                                       anakin_shape);
+    anakin_output->share_from(tmp_anakin_tensor);
+  }
+  net_->prediction();
+}
+
+template <typename TargetT, Precision PrecisionType, OpRunType RunType>
+void AnakinEngine<TargetT, PrecisionType, RunType>::Freeze() {
+  PADDLE_ENFORCE(graph_->Freeze(), "Freeze anakin subgraph.");
+}
+
+template <typename TargetT, Precision PrecisionType, OpRunType RunType>
+void AnakinEngine<TargetT, PrecisionType, RunType>::Optimize() {
+  PADDLE_ENFORCE(graph_->Optimize(), "Graph optimization.");
+}
+
+template <typename TargetT, Precision PrecisionType, OpRunType RunType>
+std::unique_ptr<AnakinEngine<TargetT, PrecisionType, RunType>>
+AnakinEngine<TargetT, PrecisionType, RunType>::Clone() {
+  auto *engine = new AnakinEngine();
+  engine->net_ = std::move(net_->Clone());
+  return std::unique_ptr<AnakinEngine>(engine);
+}
+
+template class AnakinEngine<::anakin::saber::NV, ::anakin::Precision::FP32>;
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/anakin/engine.h

+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <algorithm>
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+#include "paddle/fluid/framework/lod_tensor.h"
+#include "paddle/fluid/inference/engine.h"
+#include "paddle/fluid/inference/utils/singleton.h"
+
+#include "framework/core/net/net.h"
+#include "framework/core/types.h"
+#include "framework/graph/graph.h"
+#include "saber/saber_types.h"
+
+namespace anakin {
+
+template <typename, Precision, OpRunType>
+class Net;
+
+namespace graph {
+template <typename, Precision>
+class Graph;
+}  // namespace graph
+}  // namespace anakin
+
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+template <typename TargetT, ::anakin::Precision PrecisionType,
+          ::anakin::OpRunType RunType = ::anakin::OpRunType::ASYNC>
+class AnakinEngine {
+ public:
+  explicit AnakinEngine(bool need_summary = false);
+  ~AnakinEngine();
+  void InitGraph();
+  void SetInputShape(const std::string &name, std::vector<int> shape);
+  void AddOp(const std::string &name, const std::string &type,
+             const std::vector<std::string> &inputs,
+             const std::vector<std::string> &outputs);
+
+  template <typename T>
+  void AddOpAttr(const std::string &op_name, const std::string &attr_name,
+                 const T &attr_value) {
+    PADDLE_ENFORCE(graph_->AddOpAttr(op_name, attr_name, attr_value),
+                   "Add operation's attribution.");
+  }
+
+  std::unique_ptr<AnakinEngine> Clone();
+  void Freeze();
+  void Optimize();
+  void Execute(const std::map<std::string, framework::LoDTensor *> &inputs,
+               const std::map<std::string, framework::LoDTensor *> &outputs);
+
+ private:
+  using NetT = ::anakin::Net<TargetT, PrecisionType, RunType>;
+  using GraphT = ::anakin::graph::Graph<TargetT, PrecisionType>;
+  std::unique_ptr<GraphT> graph_;
+  std::unique_ptr<NetT> net_;
+};
+
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/anakin/test_anakin_engine.cc

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <map>
+
+#include "framework/core/net/net.h"
+#include "framework/graph/graph.h"
+#include "framework/graph/graph_global_mem.h"
+#include "paddle/fluid/inference/anakin/engine.h"
+
+using anakin::graph::GraphGlobalMem;
+using anakin::AK_FLOAT;
+using anakin::Precision;
+using anakin::saber::NV;
+using anakin::saber::X86;
+using anakin::saber::Shape;
+using anakin::PBlock;
+using anakin::PTuple;
+namespace paddle {
+namespace inference {
+namespace anakin {
+
+class TestAnakinEngine : public ::testing::Test {
+ protected:
+  void SetUp() override;
+  void TearDown() override {}
+
+ protected:
+  using AnakinNvEngineT = AnakinEngine<NV, Precision::FP32>;
+  std::unique_ptr<AnakinNvEngineT> engine_{nullptr};
+};
+
+void TestAnakinEngine::SetUp() {
+  engine_.reset(new AnakinEngine<NV, Precision::FP32>(true));
+}
+
+TEST_F(TestAnakinEngine, Execute) {
+  engine_->AddOp("op1", "Dense", {"x"}, {"y"});
+  engine_->AddOpAttr("op1", "out_dim", 2);
+  engine_->AddOpAttr("op1", "bias_term", false);
+  engine_->AddOpAttr("op1", "axis", 1);
+  std::vector<int> shape = {1, 1, 1, 2};
+  Shape tmp_shape(shape);
+  // PBlock<NV> weight1(tmp_shape);
+  auto *weight1 =
+      GraphGlobalMem<NV>::Global().template new_block<AK_FLOAT>(tmp_shape);
+  // auto *weight1 = new PBlock<NV>(tmp_shape, AK_FLOAT);
+
+  float *cpu_data = static_cast<float *>(weight1->h_tensor().mutable_data());
+  cpu_data[0] = 2.;
+  weight1->d_tensor().set_shape(tmp_shape);
+  weight1->d_tensor().copy_from(weight1->h_tensor());
+  engine_->AddOpAttr("op1", "weight_1", *weight1);
+
+  engine_->Freeze();
+  // PTuple<int> input_shape = {1};
+  // engine_->AddOpAttr("x", "input_shape", input_shape);
+  engine_->SetInputShape("x", {1, 1, 1, 1});
+  engine_->Optimize();
+  engine_->InitGraph();
+  framework::LoDTensor x;
+  framework::LoDTensor y;
+  x.Resize({1, 1, 1, 1});
+  y.Resize({1, 1, 1, 2});
+  auto *x_data = x.mutable_data<float>(platform::CUDAPlace());
+  float x_data_cpu[] = {1.};
+  cudaMemcpy(x_data, x_data_cpu, sizeof(float), cudaMemcpyHostToDevice);
+
+  std::map<std::string, framework::LoDTensor *> inputs = {{"x", &x}};
+  auto *y_data = y.mutable_data<float>(platform::CUDAPlace());
+  std::map<std::string, framework::LoDTensor *> outputs = {{"y", &y}};
+
+  engine_->Execute(inputs, outputs);
+  auto *y_data_gpu = y_data;
+  float y_data_cpu[2];
+  cudaMemcpy(y_data_cpu, y_data_gpu, sizeof(float) * 2, cudaMemcpyDeviceToHost);
+  LOG(INFO) << "output value: " << y_data_cpu[0] << ", " << y_data_cpu[1];
+}
+
+}  // namespace anakin
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/api/paddle_pass_builder.cc

@@ -13,7 +13,9 @@
 // limitations under the License.
 
 #include "paddle/fluid/inference/api/paddle_pass_builder.h"
-
+#ifdef PADDLE_WITH_CUDA
+#include <cudnn.h>
+#endif
 #include <glog/logging.h>
 
 namespace paddle {

paddle/fluid/operators/activation_cudnn.cu.cc

+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/operators/activation_op.h"
+#include "paddle/fluid/platform/cudnn_desc.h"
+
+namespace paddle {
+namespace operators {
+using framework::Tensor;
+using platform::ActivationDescriptor;
+using platform::TensorDescriptor;
+
+template <typename Functor>
+class CudnnActivationKernel
+    : public framework::OpKernel<Functor::ElEWISE_TYPE> {
+ public:
+  void Compute(const framework::ExecutionContext& context) const override {
+    framework::Tensor *X, *Out;
+    ExtractActivationTensor(context, X, Out);
+    ActivationDescriptor act_desc;
+    TensorDescriptor x_desc, out_desc;
+    x_desc.set(detail::Ref(X));
+    out_desc.set(detail::Ref(Out));
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle

paddle/fluid/operators/activation_cudnn_op.cu.cc

+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/operators/activation_op.h"
+#include "paddle/fluid/platform/cudnn_desc.h"
+
+namespace paddle {
+namespace operators {
+using framework::Tensor;
+using platform::ActivationDescriptor;
+using platform::TensorDescriptor;
+using platform::CUDADeviceContext;
+
+template <typename T>
+struct CudnnActivationFunctor {
+  using ELEMENT_TYPE = T;
+  CudnnActivationFunctor(const CUDADeviceContext& ctx, const T& c,
+                         const cudnnActivationMode_t& m)
+      : ctx_(ctx), coef_(c), mode_(m) {}
+  void operator()(const Tensor& x, Tensor* out) {
+    ActivationDescriptor act_desc;
+    act_desc.set(mode_, coef_);
+    TensorDescriptor x_desc, out_desc;
+    x_desc.set(x);
+    out_desc.set(detail::Ref(out));
+    PADDLE_ENFORCE(platform::dynload::cudnnActivationForward(
+        ctx_.cudnn_handle(), act_desc.desc(),
+        platform::CudnnDataType<T>::kOne(), x_desc.desc(), x.data<T>(),
+        platform::CudnnDataType<T>::kZero(), out_desc.desc(),
+        out->mutable_data<T>(ctx_.GetPlace())));
+  }
+  const CUDADeviceContext& ctx_;
+  const T coef_;
+  const cudnnActivationMode_t mode_;
+};
+
+template <typename T>
+struct CudnnActivationGradFunctor {
+  using ELEMENT_TYPE = T;
+  CudnnActivationGradFunctor(const CUDADeviceContext& ctx, const T& c,
+                             const cudnnActivationMode_t& m)
+      : ctx_(ctx), coef_(c), mode_(m) {}
+  void operator()(const Tensor& x, const Tensor& out, const Tensor dout,
+                  Tensor* dx) {
+    ActivationDescriptor act_desc;
+    act_desc.set(mode_, coef_);
+    TensorDescriptor x_desc, out_desc, dout_desc, dx_desc;
+    x_desc.set(x);
+    out_desc.set(out);
+    dout_desc.set(dout);
+    dx_desc.set(detail::Ref(dx));
+    PADDLE_ENFORCE(platform::dynload::cudnnActivationBackward(
+        ctx_.cudnn_handle(), act_desc.desc(),
+        platform::CudnnDataType<T>::kOne(), out_desc.desc(), out.data<T>(),
+        dout_desc.desc(), dout.data<T>(), x_desc.desc(), x.data<T>(),
+        platform::CudnnDataType<T>::kZero(), dx_desc.desc(),
+        dx->mutable_data<T>(ctx_.GetPlace())));
+  }
+  const CUDADeviceContext& ctx_;
+  const T coef_;
+  const cudnnActivationMode_t mode_;
+};
+
+template <typename T>
+struct CudnnReluFunctor : public CudnnActivationFunctor<T> {
+  explicit CudnnReluFunctor(const CUDADeviceContext& ctx)
+      : CudnnActivationFunctor<T>(ctx, 0.0, CUDNN_ACTIVATION_RELU) {}
+};
+template <typename T>
+struct CudnnReluGradFunctor : public CudnnActivationGradFunctor<T> {
+  explicit CudnnReluGradFunctor(const CUDADeviceContext& ctx)
+      : CudnnActivationGradFunctor<T>(ctx, 0.0, CUDNN_ACTIVATION_RELU) {}
+};
+
+template <typename T>
+struct CudnnRelu6Functor : public CudnnActivationFunctor<T> {
+  explicit CudnnRelu6Functor(const CUDADeviceContext& ctx)
+      : CudnnActivationFunctor<T>(ctx, 6.0, CUDNN_ACTIVATION_CLIPPED_RELU) {}
+};
+template <typename T>
+struct CudnnRelu6GradFunctor : public CudnnActivationGradFunctor<T> {
+  explicit CudnnRelu6GradFunctor(const CUDADeviceContext& ctx)
+      : CudnnActivationGradFunctor<T>(ctx, 6.0, CUDNN_ACTIVATION_CLIPPED_RELU) {
+  }
+};
+
+template <typename T>
+struct CudnnSigmoidFunctor : public CudnnActivationFunctor<T> {
+  explicit CudnnSigmoidFunctor(const CUDADeviceContext& ctx)
+      : CudnnActivationFunctor<T>(ctx, 0.0, CUDNN_ACTIVATION_SIGMOID) {}
+};
+template <typename T>
+struct CudnnSigmoidGradFunctor : public CudnnActivationGradFunctor<T> {
+  explicit CudnnSigmoidGradFunctor(const CUDADeviceContext& ctx)
+      : CudnnActivationGradFunctor<T>(ctx, 0.0, CUDNN_ACTIVATION_SIGMOID) {}
+};
+
+template <typename T>
+struct CudnnTanhFunctor : public CudnnActivationFunctor<T> {
+  explicit CudnnTanhFunctor(const CUDADeviceContext& ctx)
+      : CudnnActivationFunctor<T>(ctx, 0.0, CUDNN_ACTIVATION_TANH) {}
+};
+template <typename T>
+struct CudnnTanhGradFunctor : public CudnnActivationGradFunctor<T> {
+  explicit CudnnTanhGradFunctor(const CUDADeviceContext& ctx)
+      : CudnnActivationGradFunctor<T>(ctx, 0.0, CUDNN_ACTIVATION_TANH) {}
+};
+
+template <typename Functor>
+class CudnnActivationKernel
+    : public framework::OpKernel<typename Functor::ELEMENT_TYPE> {
+ public:
+  using T = typename Functor::ELEMENT_TYPE;
+  void Compute(const framework::ExecutionContext& context) const override {
+    const framework::Tensor* X = nullptr;
+    framework::Tensor* Out = nullptr;
+    ExtractActivationTensor(context, &X, &Out);
+    Out->mutable_data<T>(context.GetPlace());
+    auto& dev_ctx = context.template device_context<CUDADeviceContext>();
+    Functor functor(dev_ctx);
+    functor(detail::Ref(X), Out);
+  }
+};
+
+template <typename Functor>
+class CudnnActivationGradKernel
+    : public framework::OpKernel<typename Functor::ELEMENT_TYPE> {
+ public:
+  using T = typename Functor::ELEMENT_TYPE;
+  void Compute(const framework::ExecutionContext& context) const override {
+    const framework::Tensor *X, *Out, *dOut;
+    X = Out = dOut = nullptr;
+    framework::Tensor* dX = nullptr;
+    ExtractActivationGradTensor(context, &X, &Out, &dOut, &dX);
+    dX->mutable_data<T>(context.GetPlace());
+    auto& dev_ctx = context.template device_context<CUDADeviceContext>();
+    Functor functor(dev_ctx);
+    functor(detail::Ref(X), detail::Ref(Out), detail::Ref(dOut), dX);
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace plat = paddle::platform;
+namespace ops = paddle::operators;
+
+#define FOR_EACH_CUDNN_OP_FUNCTOR(__macro)                  \
+  __macro(relu, CudnnReluFunctor, CudnnReluGradFunctor);    \
+  __macro(relu6, CudnnRelu6Functor, CudnnRelu6GradFunctor); \
+  __macro(sigmoid, CudnnTanhFunctor, CudnnTanhGradFunctor); \
+  __macro(tanh, CudnnTanhFunctor, CudnnTanhGradFunctor)
+
+#define REGISTER_ACTIVATION_CUDNN_KERNEL(act_type, functor, grad_functor) \
+  REGISTER_OP_KERNEL(act_type, CUDNN, plat::CUDAPlace,                    \
+                     ops::CudnnActivationKernel<ops::functor<float>>,     \
+                     ops::CudnnActivationKernel<ops::functor<double>>);   \
+  REGISTER_OP_KERNEL(                                                     \
+      act_type##_grad, CUDNN, plat::CUDAPlace,                            \
+      ops::CudnnActivationGradKernel<ops::grad_functor<float>>,           \
+      ops::CudnnActivationGradKernel<ops::grad_functor<double>>);
+
+FOR_EACH_CUDNN_OP_FUNCTOR(REGISTER_ACTIVATION_CUDNN_KERNEL);

paddle/fluid/operators/activation_op.cc

@@ -16,29 +16,36 @@ limitations under the License. */
 #include <string>
 #include "paddle/fluid/operators/mkldnn/mkldnn_activation_op.h"
 #include "paddle/fluid/platform/port.h"
+#ifdef PADDLE_WITH_CUDA
+#include "paddle/fluid/platform/cudnn_helper.h"
+#endif
 
 namespace paddle {
 namespace operators {
 
 using paddle::framework::Tensor;
 
-#define REGISTER_ACTIVATION_OP_MAKER(OP_NAME, OP_COMMENT)                \
-  class OP_NAME##OpMaker                                                 \
-      : public ::paddle::framework::OpProtoAndCheckerMaker {             \
-   public:                                                               \
-    void Make() override {                                               \
-      AddInput("X", "Input of " #OP_NAME " operator");                   \
-      AddOutput("Out", "Output of " #OP_NAME " operator");               \
-      AddAttr<bool>("use_mkldnn",                                        \
-                    "(bool, default false) Only used in mkldnn kernel")  \
-          .SetDefault(false);                                            \
-      AddAttr<bool>(                                                     \
-          "is_test",                                                     \
-          "(bool, default false) Set to true for inference only, false " \
-          "for training. Some layers may run faster when this is true.") \
-          .SetDefault(false);                                            \
-      AddComment(OP_COMMENT);                                            \
-    }                                                                    \
+#define REGISTER_ACTIVATION_OP_MAKER(OP_NAME, OP_COMMENT)                    \
+  class OP_NAME##OpMaker                                                     \
+      : public ::paddle::framework::OpProtoAndCheckerMaker {                 \
+   public:                                                                   \
+    void Make() override {                                                   \
+      AddInput("X", "Input of " #OP_NAME " operator");                       \
+      AddOutput("Out", "Output of " #OP_NAME " operator");                   \
+      AddAttr<bool>("use_mkldnn",                                            \
+                    "(bool, default false) Only used in mkldnn kernel")      \
+          .SetDefault(false);                                                \
+      AddAttr<bool>("use_cudnn",                                             \
+                    "(bool, default false) Only used in cudnn kernel, need " \
+                    "install cudnn")                                         \
+          .SetDefault(false);                                                \
+      AddAttr<bool>(                                                         \
+          "is_test",                                                         \
+          "(bool, default false) Set to true for inference only, false "     \
+          "for training. Some layers may run faster when this is true.")     \
+          .SetDefault(false);                                                \
+      AddComment(OP_COMMENT);                                                \
+    }                                                                        \
   }
 
 #define REGISTER_ACTIVATION_OP_GRAD_MAKER(OP_NAME, KERNEL_TYPE)              \
@@ -67,6 +74,12 @@ framework::OpKernelType GetKernelType(const framework::ExecutionContext& ctx,
                                       const std::string& name) {
   framework::LibraryType library{framework::LibraryType::kPlain};
   framework::DataLayout layout = framework::DataLayout::kAnyLayout;
+#ifdef PADDLE_WITH_CUDA
+  auto it1 = oper.Attrs().find("use_cudnn");
+  if (it1 != oper.Attrs().end() && platform::CanCUDNNBeUsed(ctx)) {
+    library = framework::LibraryType::kCUDNN;
+  }
+#endif
 #ifdef PADDLE_WITH_MKLDNN
   auto it = oper.Attrs().find("use_mkldnn");
   if (library == framework::LibraryType::kPlain && it != oper.Attrs().end() &&

paddle/fluid/operators/activation_op.h

@@ -11,6 +11,7 @@ limitations under the License. */
 
 #pragma once
 #include <glog/logging.h>
+#include <algorithm>
 #include <string>
 #include <unordered_set>
 #include <utility>
@@ -24,6 +25,7 @@ limitations under the License. */
 #include "paddle/fluid/framework/eigen.h"
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/operators/detail/safe_ref.h"
+#include "paddle/fluid/operators/math/blas.h"
 #include "paddle/fluid/platform/float16.h"
 
 #ifdef PADDLE_WITH_MKLDNN
@@ -41,53 +43,115 @@ static std::unordered_set<std::string> InplaceOpSet = {
     "floor",   "reciprocal", "relu6", "soft_relu", "hard_sigmoid",
 };
 
+static bool IsInplace(const std::string& op) {
+  bool inplace = InplaceOpSet.count(op);
+  // for op_grad
+  const int kGradSuffixLen = 4;
+  if (op.size() > kGradSuffixLen &&
+      op.compare(op.size() - kGradSuffixLen - 1, kGradSuffixLen, "grad")) {
+    inplace =
+        InplaceOpSet.count(op.substr(0, op.size() - (kGradSuffixLen + 1)));
+  }
+  return inplace;
+}
+
 /* The following operator can be used to process SelectedRows, because the
  * output of those operator for zero is zero too.
  */
 static std::unordered_set<std::string> CanBeUsedBySelectedRows = {
     "abs", "abs_grad", "square", "square_grad", "sqrt", "sqrt_grad"};
 
-static bool IsInplace(std::string op) { return InplaceOpSet.count(op); }
-
-template <typename DeviceContext, typename Functor>
-class ActivationKernel
-    : public framework::OpKernel<typename Functor::ELEMENT_TYPE> {
- public:
-  using T = typename Functor::ELEMENT_TYPE;
-
-  void Compute(const framework::ExecutionContext& context) const override {
+inline void ExtractActivationTensor(const framework::ExecutionContext& context,
+                                    const framework::Tensor** X,
+                                    framework::Tensor** Out) {
+  auto x_var = context.InputVar("X");
+  auto out_var = context.OutputVar("Out");
+  PADDLE_ENFORCE(x_var != nullptr,
+                 "Cannot get input Variable X, variable name = %s",
+                 context.op().Input("X"));
+  PADDLE_ENFORCE(out_var != nullptr,
+                 "Cannot get output Variable Out, variable name = %s",
+                 context.op().Output("Out"));
+  if (CanBeUsedBySelectedRows.count(context.op().Type())) {
+    *X = paddle::framework::GetLoDTensorOrSelectedRowsValueFromVar(*x_var);
+    *Out = paddle::framework::GetMutableLoDTensorOrSelectedRowsValueFromVar(
+        out_var);
+  } else {
+    *X = context.Input<framework::Tensor>("X");
+    *Out = context.Output<framework::Tensor>("Out");
+  }
+
+  PADDLE_ENFORCE(*Out != nullptr,
+                 "Cannot get output tensor Out, variable name = %s",
+                 context.op().Output("Out"));
+}
+
+inline void ExtractActivationGradTensor(
+    const framework::ExecutionContext& context, const framework::Tensor** X,
+    const framework::Tensor** Out, const framework::Tensor** dOut,
+    framework::Tensor** dX) {
+  auto out_var = context.InputVar("Out");
+  auto out_grad_var = context.InputVar(framework::GradVarName("Out"));
+  auto x_grad_var = context.OutputVar(framework::GradVarName("X"));
+  PADDLE_ENFORCE(out_var != nullptr,
+                 "Cannot get input Variable Out, variable name = %s",
+                 context.op().Input("Out"));
+  PADDLE_ENFORCE(out_grad_var != nullptr,
+                 "Cannot get input Variable %s, variable name = %s",
+                 framework::GradVarName("Out"),
+                 context.op().Input(framework::GradVarName("Out")));
+  PADDLE_ENFORCE(x_grad_var != nullptr,
+                 "Cannot get output Variable %s, variable name = %s",
+                 framework::GradVarName("X"),
+                 context.op().Output(framework::GradVarName("X")));
+
+  if (CanBeUsedBySelectedRows.count(context.op().Type())) {
+    *Out = paddle::framework::GetLoDTensorOrSelectedRowsValueFromVar(*out_var);
+    *dOut = paddle::framework::GetLoDTensorOrSelectedRowsValueFromVar(
+        *out_grad_var);
+    *dX = paddle::framework::GetMutableLoDTensorOrSelectedRowsValueFromVar(
+        x_grad_var);
+  } else {
+    *Out = context.Input<framework::Tensor>("Out");
+    *dOut = context.Input<framework::Tensor>(framework::GradVarName("Out"));
+    *dX = context.Output<framework::Tensor>(framework::GradVarName("X"));
+  }
+  PADDLE_ENFORCE(*dX != nullptr,
+                 "Cannot get output tensor %s, variable name = %s",
+                 framework::GradVarName("X"),
+                 context.op().Output(framework::GradVarName("X")));
+
+  bool inplace = IsInplace(context.op().Type());
+  if (!inplace) {
     auto x_var = context.InputVar("X");
-    auto out_var = context.OutputVar("Out");
     PADDLE_ENFORCE(x_var != nullptr,
-                   "Cannot get input Variable X, variable name = %s",
+                   "Cannot get input tensor X, variable name = %s",
                    context.op().Input("X"));
-    PADDLE_ENFORCE(out_var != nullptr,
-                   "Cannot get output Variable Out, variable name = %s",
-                   context.op().Output("Out"));
-
-    framework::Tensor X, *Out;
-
     if (CanBeUsedBySelectedRows.count(context.op().Type())) {
-      X = detail::Ref(
-          paddle::framework::GetLoDTensorOrSelectedRowsValueFromVar(*x_var),
-          "Cannot get input Tensor X, variable name = %s",
-          context.op().Input("X"));
-      Out = paddle::framework::GetMutableLoDTensorOrSelectedRowsValueFromVar(
-          out_var);
+      *X = paddle::framework::GetLoDTensorOrSelectedRowsValueFromVar(*x_var);
     } else {
-      X = detail::Ref(context.Input<framework::Tensor>("X"),
-                      "Cannot get input Tensor X, variable name = %s",
-                      context.op().Input("X"));
-      Out = context.Output<framework::Tensor>("Out");
+      *X = context.Input<framework::Tensor>("X");
     }
+  } else {
+    VLOG(10) << " Inplace activation of Op : " << context.op().Type();
+    *X = *dX;
+  }
+}
 
-    PADDLE_ENFORCE(Out != nullptr,
-                   "Cannot get output tensor Out, variable name = %s",
-                   context.op().Output("Out"));
+template <typename DeviceContext, typename Functor>
+class ActivationKernel
+    : public framework::OpKernel<typename Functor::ELEMENT_TYPE> {
+ public:
+  using T = typename Functor::ELEMENT_TYPE;
 
+  void Compute(const framework::ExecutionContext& context) const override {
+    const framework::Tensor* X = nullptr;
+    framework::Tensor* Out = nullptr;
+    ExtractActivationTensor(context, &X, &Out);
     Out->mutable_data<T>(context.GetPlace());
-    auto x = framework::EigenVector<T>::Flatten(X);
-    auto out = framework::EigenVector<T>::Flatten(*Out);
+
+    auto x = framework::EigenVector<T>::Flatten(detail::Ref(X));
+    auto out = framework::EigenVector<T>::Flatten(detail::Ref(Out));
     auto* place =
         context.template device_context<DeviceContext>().eigen_device();
     Functor functor;
@@ -106,55 +170,15 @@ class ActivationGradKernel
  public:
   using T = typename Functor::ELEMENT_TYPE;
   void Compute(const framework::ExecutionContext& context) const override {
-    auto out_var = context.InputVar("Out");
-    auto out_grad_var = context.InputVar(framework::GradVarName("Out"));
-    auto x_grad_var = context.OutputVar(framework::GradVarName("X"));
-    PADDLE_ENFORCE(out_var != nullptr,
-                   "Cannot get input Variable Out, variable name = %s",
-                   context.op().Input("Out"));
-    PADDLE_ENFORCE(out_grad_var != nullptr,
-                   "Cannot get input Variable %s, variable name = %s",
-                   framework::GradVarName("Out"),
-                   context.op().Input(framework::GradVarName("Out")));
-    PADDLE_ENFORCE(x_grad_var != nullptr,
-                   "Cannot get output Variable %s, variable name = %s",
-                   framework::GradVarName("X"),
-                   context.op().Output(framework::GradVarName("X")));
-
-    framework::Tensor Out, dOut, *dX;
-    if (CanBeUsedBySelectedRows.count(context.op().Type())) {
-      Out = detail::Ref(
-          paddle::framework::GetLoDTensorOrSelectedRowsValueFromVar(*out_var),
-          "Cannot get input Tensor Out, variable name = %s",
-          context.op().Input("Out"));
-      dOut =
-          detail::Ref(paddle::framework::GetLoDTensorOrSelectedRowsValueFromVar(
-                          *out_grad_var),
-                      "Cannot get input Tensor %s, variable name = %s",
-                      framework::GradVarName("Out"),
-                      context.op().Input(framework::GradVarName("Out")));
-      dX = paddle::framework::GetMutableLoDTensorOrSelectedRowsValueFromVar(
-          x_grad_var);
-    } else {
-      Out = detail::Ref(context.Input<framework::Tensor>("Out"),
-                        "Cannot get input Tensor Out, variable name = %s",
-                        context.op().Input("Out"));
-      dOut = detail::Ref(
-          context.Input<framework::Tensor>(framework::GradVarName("Out")),
-          "Cannot get input Tensor %s, variable name = %s",
-          framework::GradVarName("Out"),
-          context.op().Input(framework::GradVarName("Out")));
-      dX = context.Output<framework::Tensor>(framework::GradVarName("X"));
-    }
-    PADDLE_ENFORCE(dX != nullptr,
-                   "Cannot get output tensor %s, variable name = %s",
-                   framework::GradVarName("X"),
-                   context.op().Output(framework::GradVarName("X")));
+    const framework::Tensor *X, *Out, *dOut;
+    framework::Tensor* dX = nullptr;
+    X = Out = dOut = nullptr;
+    ExtractActivationGradTensor(context, &X, &Out, &dOut, &dX);
     dX->mutable_data<T>(context.GetPlace());
-
-    auto dout = framework::EigenVector<T>::Flatten(dOut);
-    auto out = framework::EigenVector<T>::Flatten(Out);
-    auto dx = framework::EigenVector<T>::Flatten(*dX);
+    auto dout = framework::EigenVector<T>::Flatten(detail::Ref(dOut));
+    auto out = framework::EigenVector<T>::Flatten(detail::Ref(Out));
+    auto dx = framework::EigenVector<T>::Flatten(detail::Ref(dX));
+    auto x = framework::EigenVector<T>::Flatten(detail::Ref(X));
     auto* place =
         context.template device_context<DeviceContext>().eigen_device();
     Functor functor;
@@ -162,27 +186,7 @@ class ActivationGradKernel
     for (auto& attr : attrs) {
       *attr.second = context.Attr<float>(attr.first);
     }
-    bool inplace = functor.Inplace();
-    if (!inplace) {
-      auto x_var = context.InputVar("X");
-      PADDLE_ENFORCE(x_var != nullptr,
-                     "Cannot get input tensor X, variable name = %s",
-                     context.op().Input("X"));
-      framework::Tensor X;
-      if (CanBeUsedBySelectedRows.count(context.op().Type())) {
-        X = detail::Ref(
-            paddle::framework::GetLoDTensorOrSelectedRowsValueFromVar(*x_var));
-      } else {
-        X = detail::Ref(context.Input<framework::Tensor>("X"));
-      }
-
-      auto x = framework::EigenVector<T>::Flatten(X);
-      functor(*place, x, out, dout, dx);
-    } else {
-      VLOG(10) << " Inplace activation ";
-      auto x = framework::EigenVector<T>::Flatten(*dX);
-      functor(*place, x, out, dout, dx);
-    }
+    functor(*place, x, out, dout, dx);
   }
 };
 
@@ -214,7 +218,6 @@ struct SigmoidFunctor : public BaseActivationFunctor<T> {
 
 template <typename T>
 struct SigmoidGradFunctor : public BaseActivationFunctor<T> {
-  bool Inplace() const { return IsInplace("sigmoid"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {
@@ -269,7 +272,6 @@ struct ExpFunctor : public BaseActivationFunctor<T> {
 
 template <typename T>
 struct ExpGradFunctor : public BaseActivationFunctor<T> {
-  bool Inplace() const { return IsInplace("exp"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {
@@ -288,7 +290,6 @@ struct ReluFunctor : public BaseActivationFunctor<T> {
 
 template <typename T>
 struct ReluGradFunctor : public BaseActivationFunctor<T> {
-  bool Inplace() const { return IsInplace("relu"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {
@@ -301,8 +302,28 @@ template <typename T>
 struct GeluFunctor : public BaseActivationFunctor<T> {
   template <typename Device, typename X, typename Out>
   void operator()(Device d, X x, Out out) const {
+// Because the execute or device context can not be deliver here, it keep the
+// marco for NVCC.
+#if defined(PADDLE_WITH_MKLML) && !defined(_WIN32) && !defined(__APPLE__) && \
+    !defined(__OSX__) && !defined(PADDLE_WITH_CUDA)
+    auto x_data = x.data();
+    auto out_data = out.data();
+    int n = std::min(x.size(), out.size());
+
+    std::memset(out_data, 0, n * sizeof(T));
+    math::CBlas<T>::AXPY(n, static_cast<T>(M_SQRT1_2), x_data, 1, out_data, 1);
+    math::CBlas<T>::VMERF(n, out_data, out_data, VML_LA);
+    for (int i = 0; i < n; i++) {
+      out_data[i] += static_cast<T>(1);
+    }
+    math::CBlas<T>::VMUL(n, x_data, out_data, out_data);
+    for (int i = 0; i < n; i++) {
+      out_data[i] *= static_cast<T>(0.5);
+    }
+#else
     auto temp = (x * static_cast<T>(M_SQRT1_2)).erf();
     out.device(d) = x * static_cast<T>(0.5) * (static_cast<T>(1) + temp);
+#endif
   }
 };
 
@@ -331,7 +352,6 @@ struct TanhFunctor : public BaseActivationFunctor<T> {
 
 template <typename T>
 struct TanhGradFunctor : public BaseActivationFunctor<T> {
-  bool Inplace() const { return IsInplace("tanh"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {
@@ -437,7 +457,6 @@ struct SqrtFunctor : public BaseActivationFunctor<T> {
 
 template <typename T>
 struct SqrtGradFunctor : public BaseActivationFunctor<T> {
-  bool Inplace() const { return IsInplace("sqrt"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {
@@ -456,7 +475,6 @@ struct CeilFunctor : public BaseActivationFunctor<T> {
 
 template <typename T>
 struct ZeroGradFunctor : public BaseActivationFunctor<T> {
-  bool Inplace() const { return IsInplace("ceil"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {
@@ -573,7 +591,6 @@ struct ReciprocalFunctor : public BaseActivationFunctor<T> {
 
 template <typename T>
 struct ReciprocalGradFunctor : public BaseActivationFunctor<T> {
-  bool Inplace() const { return IsInplace("reciprocal"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {
@@ -673,7 +690,6 @@ struct Relu6GradFunctor : public BaseActivationFunctor<T> {
   typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
     return {{"threshold", &threshold}};
   }
-  bool Inplace() const { return IsInplace("relu6"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {
@@ -755,7 +771,6 @@ struct SoftReluGradFunctor : public BaseActivationFunctor<T> {
   typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
     return {{"threshold", &threshold}};
   }
-  bool Inplace() const { return IsInplace("soft_relu"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {
@@ -936,7 +951,6 @@ struct HardSigmoidGradFunctor : public BaseActivationFunctor<T> {
   typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
     return {{"slope", &slope}, {"offset", &offset}};
   }
-  bool Inplace() { return IsInplace("hard_sigmoid"); }
   template <typename Device, typename X, typename Out, typename dOut,
             typename dX>
   void operator()(Device d, X x, Out out, dOut dout, dX dx) const {

paddle/fluid/operators/alloc_continuous_space_op.cc

+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <vector>
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/framework/var_type.h"
+#include "paddle/fluid/operators/math/math_function.h"
+
+namespace paddle {
+namespace operators {
+
+static framework::proto::VarType::Type kDefaultDtype =
+    framework::proto::VarType::Type::VarType_Type_BOOL;
+
+template <typename DeviceContext, typename T>
+class AllocContinuousSpaceKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext &context) const override {
+    auto &in_var_names = context.Inputs("Input");
+    auto &out_var_names = context.Outputs("Output");
+    auto &in_vars = context.MultiInputVar("Input");
+    auto out_vars = context.MultiOutputVar("Output");
+
+    PADDLE_ENFORCE_GT(in_var_names.size(), static_cast<size_t>(0));
+    PADDLE_ENFORCE_EQ(in_var_names.size(), out_var_names.size());
+
+    for (size_t i = 0; i < in_var_names.size(); ++i) {
+      // Only support LoDTensor
+      PADDLE_ENFORCE_NOT_NULL(in_vars[i], "%s should not be nullptr,",
+                              in_var_names[i]);
+      PADDLE_ENFORCE_NOT_NULL(out_vars[i], "%s should not be nullptr,",
+                              out_var_names[i]);
+      PADDLE_ENFORCE(in_vars[i]->IsType<framework::LoDTensor>());
+      PADDLE_ENFORCE(out_vars[i]->IsType<framework::LoDTensor>());
+    }
+
+    auto in_tensors = context.MultiInput<framework::LoDTensor>("Input");
+
+    if (context.Attr<bool>("check_name")) {
+      for (size_t i = 0; i < in_var_names.size(); ++i) {
+        PADDLE_ENFORCE_EQ(in_var_names[i], out_var_names[i]);
+      }
+    } else {
+      // Init the output as input
+      for (size_t i = 0; i < in_tensors.size(); ++i) {
+        out_vars[i]->GetMutable<framework::LoDTensor>()->Resize(
+            in_tensors[i]->dims());
+      }
+    }
+
+    auto &dev_ctx = context.template device_context<DeviceContext>();
+
+    // Get numel and dtype
+    size_t numel = 0;
+    auto dtype = kDefaultDtype;
+    GetMemSizeAndDtype(in_tensors, in_var_names, &numel, &dtype);
+
+    // Alloc the continuous space
+    auto fused_tensor = context.Output<framework::LoDTensor>("FusedOutput");
+    fused_tensor->Resize(framework::make_ddim({static_cast<int64_t>(numel)}))
+        .mutable_data(context.GetPlace(), dtype);
+
+    // Init the continuous space
+    auto out_tensors = context.MultiOutput<framework::LoDTensor>("Output");
+    int64_t offset = 0;
+    if (context.Attr<bool>("copy_data")) {
+      for (size_t i = 0; i < in_var_names.size(); ++i) {
+        int64_t len = out_tensors[i]->numel();
+        auto sub_tensor = fused_tensor->Slice(offset, offset + len);
+        offset += len;
+        framework::TensorCopy(*out_tensors[i], context.GetPlace(), dev_ctx,
+                              &sub_tensor);
+      }
+    } else if (context.Attr<bool>("set_constant")) {
+      math::SetConstant<DeviceContext, T> set_constant;
+      set_constant(dev_ctx, fused_tensor,
+                   static_cast<T>(context.Attr<float>("constant")));
+    }
+
+    // Make the outputs point to the continuous space.
+    offset = 0;
+    for (size_t i = 0; i < out_tensors.size(); ++i) {
+      int64_t len = out_tensors[i]->numel();
+      auto dim = out_tensors[i]->dims();
+      out_tensors[i]
+          ->ShareDataWith(fused_tensor->Slice(offset, offset + len))
+          .Resize(dim);
+      offset += len;
+      VLOG(10) << "alloc_space_for_vars: output(" << out_var_names[i]
+               << ") ,dim:(" << dim << ")"
+               << " Address: " << out_tensors[i]->data<void>();
+    }
+  }
+
+  void GetMemSizeAndDtype(
+      const std::vector<const framework::LoDTensor *> &lod_tensors,
+      const std::vector<std::string> var_names, size_t *numel,
+      framework::proto::VarType::Type *dtype) const {
+    PADDLE_ENFORCE_EQ(lod_tensors.size(), var_names.size());
+    *numel = 0;
+    for (size_t i = 0; i < var_names.size(); ++i) {
+      PADDLE_ENFORCE(lod_tensors[i]->IsInitialized(), "%s is not initialized.",
+                     var_names[i]);
+
+      auto p_dtype = lod_tensors[i]->type();
+      if (*dtype == kDefaultDtype) {
+        PADDLE_ENFORCE_NE(p_dtype, kDefaultDtype, "%s's type should not be %s.",
+                          var_names[i], kDefaultDtype);
+        *dtype = p_dtype;
+      }
+      PADDLE_ENFORCE_EQ(p_dtype, *dtype, "Input vars is not equal.");
+
+      auto size = lod_tensors[i]->numel();
+      PADDLE_ENFORCE_GT(size, 0);
+      VLOG(10) << "alloc_space_for_vars: input(" << var_names[i] << ") ,dim:("
+               << lod_tensors[i]->dims() << ")";
+      *numel += size;
+    }
+  }
+};
+
+class AllocContinuousSpaceOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+  void InferShape(framework::InferShapeContext *ctx) const override {}
+};
+
+class AllocContinuousSpaceOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("Input",
+             "(vector<LoDTensor>) The input tensors of"
+             " alloc_continuous_space operator.")
+        .AsDuplicable();
+    AddOutput("Output",
+              "(vector<LoDTensor>) The output "
+              "tensors of alloc_continuous_space operator. And the address "
+              "of output tensors are continuous, they are sliced from the "
+              "tensor of FusedOutput.")
+        .AsDuplicable();
+    AddOutput("FusedOutput",
+              "(LoDTensor) The output tensor "
+              "of alloc_continuous_space operator. And the tensors of"
+              " Output is sliced from the tensor of FusedOutput.");
+    AddAttr<bool>("copy_data", "Whether to copy the Input value to Output.")
+        .SetDefault(false);
+    AddAttr<bool>("set_constant",
+                  "Whether to set the Output with a constant value.")
+        .SetDefault(false);
+    AddAttr<float>("constant",
+                   "If set_constant is true, the constant value will be used "
+                   "to set the Output.")
+        .SetDefault(0.0);
+    AddAttr<bool>("check_name",
+                  "Whether to check the name of Input and Output to ensure "
+                  "they are the same separately.")
+        .SetDefault(false);
+    AddComment(R"DOC(
+AllocContinuousSpace Operator.
+
+alloc_continuous_space is used to make the address of Output
+continuous according to the Input. This Op will alloc a big tensor
+according to the tensors of Input, the dtype is the same with those input tensors,
+the size is the sum of those input tensors' numel, and the dim of the big
+tensor is {sum(numel)}. And the big tensor is stored in FusedOutput.
+The tensors of Output are sliced from the tensor of FusedOutput.
+Note that, the dtype of Input should be the same, and the dim of Input
+and Output should equal.
+The tensors of Input and Output could be the same or different. And
+alloc_continuous_space allows copying the value of Input to Output, or
+setting the Output with a constant value.
+
+)DOC");
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+REGISTER_OPERATOR(alloc_continuous_space,
+                  paddle::operators::AllocContinuousSpaceOp,
+                  paddle::operators::AllocContinuousSpaceOpMaker);
+namespace ops = paddle::operators;
+REGISTER_OP_CPU_KERNEL(
+    alloc_continuous_space,
+    ops::AllocContinuousSpaceKernel<paddle::platform::CPUDeviceContext, int>,
+    ops::AllocContinuousSpaceKernel<paddle::platform::CPUDeviceContext, float>,
+    ops::AllocContinuousSpaceKernel<paddle::platform::CPUDeviceContext,
+                                    double>);
+
+#ifdef PADDLE_WITH_CUDA
+REGISTER_OP_CUDA_KERNEL(
+    alloc_continuous_space,
+    ops::AllocContinuousSpaceKernel<paddle::platform::CUDADeviceContext, int>,
+    ops::AllocContinuousSpaceKernel<paddle::platform::CUDADeviceContext, float>,
+    ops::AllocContinuousSpaceKernel<paddle::platform::CUDADeviceContext,
+                                    double>);
+#endif

paddle/fluid/operators/benchmark/op_tester.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/benchmark/op_tester.h"
+#include <fstream>
 #include "gflags/gflags.h"
 #include "gtest/gtest.h"
 #include "paddle/fluid/framework/op_info.h"
@@ -28,6 +29,7 @@ namespace operators {
 namespace benchmark {
 
 DEFINE_string(op_config_list, "", "Path of op config file.");
+DEFINE_int32(specified_config_id, -1, "Test the specified op config.");
 
 void OpTester::Init(const std::string &filename) {
   Init(OpTesterConfig(filename));
@@ -147,7 +149,7 @@ void OpTester::CreateInputVarDesc() {
     var->SetShape(input->dims);
 
     op_desc_.SetInput(name, {var_name});
-    inputs_.push_back(var_name);
+    input_lods_[var_name] = input->lod;
   }
 }
 
@@ -162,7 +164,6 @@ void OpTester::CreateOutputVarDesc() {
     var->SetDataType(framework::proto::VarType::FP32);
 
     op_desc_.SetOutput(name, {var_name});
-    outputs_.push_back(var_name);
   }
 }
 
@@ -218,16 +219,26 @@ void OpTester::CreateVariables(framework::Scope *scope) {
     }
   }
 
-  // Allocate memory for input tensor
-  for (auto &name : inputs_) {
-    VLOG(3) << "Allocate memory for tensor " << name;
-    auto &var_desc = vars_[name];
+  for (auto &item : input_lods_) {
+    // Allocate memory for input tensor
+    auto &var_name = item.first;
+    VLOG(3) << "Allocate memory for tensor " << var_name;
+
+    auto &var_desc = vars_[var_name];
     std::vector<int64_t> shape = var_desc->GetShape();
 
-    auto *var = scope->Var(name);
+    auto *var = scope->Var(var_name);
     auto *tensor = var->GetMutable<framework::LoDTensor>();
     SetupTensor<float>(tensor, shape, static_cast<float>(0.0),
                        static_cast<float>(1.0));
+
+    VLOG(3) << "Set lod for tensor " << var_name;
+    std::vector<std::vector<size_t>> &lod_vec = item.second;
+    framework::LoD lod;
+    for (size_t i = 0; i < lod_vec.size(); ++i) {
+      lod.push_back(lod_vec[i]);
+    }
+    tensor->set_lod(lod);
   }
 }
 
@@ -282,10 +293,32 @@ std::string OpTester::DebugString() {
 }
 
 TEST(op_tester, base) {
-  OpTester tester;
   if (!FLAGS_op_config_list.empty()) {
-    tester.Init(FLAGS_op_config_list);
+    std::ifstream fin(FLAGS_op_config_list, std::ios::in | std::ios::binary);
+    PADDLE_ENFORCE(static_cast<bool>(fin), "Cannot open file %s",
+                   FLAGS_op_config_list.c_str());
+    std::vector<OpTesterConfig> op_configs;
+    while (!fin.eof()) {
+      OpTesterConfig config;
+      bool result = config.Init(fin);
+      if (result) {
+        op_configs.push_back(config);
+      }
+    }
+    if (FLAGS_specified_config_id >= 0 &&
+        FLAGS_specified_config_id < static_cast<int>(op_configs.size())) {
+      OpTester tester;
+      tester.Init(op_configs[FLAGS_specified_config_id]);
+      tester.Run();
+    } else {
+      for (size_t i = 0; i < op_configs.size(); ++i) {
+        OpTester tester;
+        tester.Init(op_configs[i]);
+        tester.Run();
+      }
+    }
   } else {
+    OpTester tester;
     OpTesterConfig config;
     config.op_type = "elementwise_add";
     config.inputs.resize(2);
@@ -294,8 +327,8 @@ TEST(op_tester, base) {
     config.inputs[1].name = "Y";
     config.inputs[1].dims = {64, 1};
     tester.Init(config);
+    tester.Run();
   }
-  tester.Run();
 }
 
 }  // namespace benchmark

paddle/fluid/operators/benchmark/op_tester.h

@@ -57,8 +57,7 @@ class OpTester {
   std::string type_;
   framework::OpDesc op_desc_;
   std::unordered_map<std::string, std::unique_ptr<framework::VarDesc>> vars_;
-  std::vector<std::string> inputs_;
-  std::vector<std::string> outputs_;
+  std::unordered_map<std::string, std::vector<std::vector<size_t>>> input_lods_;
   std::unique_ptr<framework::OperatorBase> op_;
   platform::Place place_;
   std::unique_ptr<framework::Scope> scope_;

paddle/fluid/operators/benchmark/op_tester_config.cc

@@ -33,21 +33,64 @@ static bool EndWith(const std::string& str, const std::string& substr) {
   return str.rfind(substr) == (str.length() - substr.length());
 }
 
-static void EraseEndSep(std::string* str) {
-  std::string substr = kSepBetweenItems;
+static void EraseEndSep(std::string* str,
+                        std::string substr = kSepBetweenItems) {
   if (EndWith(*str, substr)) {
     str->erase(str->length() - substr.length(), str->length());
   }
 }
 
-static std::vector<int64_t> ParseDims(std::string dims_str) {
-  std::vector<int64_t> dims;
+void OpInputConfig::ParseDims(std::istream& is) {
+  std::string dims_str;
+  is >> dims_str;
+
+  dims.clear();
   std::string token;
   std::istringstream token_stream(dims_str);
   while (std::getline(token_stream, token, 'x')) {
     dims.push_back(std::stoi(token));
   }
-  return dims;
+}
+
+void OpInputConfig::ParseLoD(std::istream& is) {
+  std::string lod_str;
+  std::string start_sep =
+      std::string(kStartSeparator) + std::string(kStartSeparator);
+  std::string end_sep = std::string(kEndSeparator) + std::string(kEndSeparator);
+
+  std::string sep;
+  is >> sep;
+  if (StartWith(sep, start_sep)) {
+    lod_str += sep;
+    while (!EndWith(sep, end_sep)) {
+      is >> sep;
+      lod_str += sep;
+    }
+  }
+  EraseEndSep(&lod_str);
+  PADDLE_ENFORCE_GE(lod_str.length(), 4U);
+  VLOG(4) << "lod: " << lod_str << ", length: " << lod_str.length();
+
+  // Parse the lod_str
+  lod.clear();
+  for (size_t i = 1; i < lod_str.length() - 1;) {
+    if (lod_str[i] == '{') {
+      std::vector<size_t> level;
+      while (lod_str[i] != '}') {
+        ++i;
+
+        std::string number;
+        while (lod_str[i] >= '0' && lod_str[i] <= '9') {
+          number += lod_str[i];
+          ++i;
+        }
+        level.push_back(atoi(number.c_str()));
+      }
+      lod.push_back(level);
+    } else if (lod_str[i] == '}') {
+      ++i;
+    }
+  }
 }
 
 OpInputConfig::OpInputConfig(std::istream& is) {
@@ -60,9 +103,9 @@ OpInputConfig::OpInputConfig(std::istream& is) {
         is >> name;
         EraseEndSep(&name);
       } else if (sep == "dims" || sep == "dims:") {
-        std::string dims_str;
-        is >> dims_str;
-        dims = ParseDims(dims_str);
+        ParseDims(is);
+      } else if (sep == "lod" || sep == "lod:") {
+        ParseLoD(is);
       }
     }
   }
@@ -76,7 +119,7 @@ OpTesterConfig::OpTesterConfig(const std::string& filename) {
   Init(fin);
 }
 
-void OpTesterConfig::Init(std::istream& is) {
+bool OpTesterConfig::Init(std::istream& is) {
   std::string sep;
   is >> sep;
   if (sep == kStartSeparator) {
@@ -95,9 +138,40 @@ void OpTesterConfig::Init(std::istream& is) {
       } else if (sep == "input" || sep == "input:") {
         OpInputConfig input_config(is);
         inputs.push_back(input_config);
+      } else if (sep == "attrs" || sep == "attrs:") {
+        ParseAttrs(is);
+      } else {
+        if (sep != kEndSeparator) {
+          return false;
+        }
       }
     }
+  } else {
+    return false;
+  }
+  return true;
+}
+
+bool OpTesterConfig::ParseAttrs(std::istream& is) {
+  std::string sep;
+  is >> sep;
+  if (sep == kStartSeparator) {
+    while (true) {
+      std::string key;
+      is >> key;
+      if (key == kEndSeparator) {
+        break;
+      }
+
+      std::string value;
+      is >> value;
+      EraseEndSep(&key, ":");
+      EraseEndSep(&value);
+
+      attrs[key] = value;
+    }
   }
+  return true;
 }
 
 const OpInputConfig* OpTesterConfig::GetInput(const std::string& name) {

paddle/fluid/operators/benchmark/op_tester_config.h

@@ -16,6 +16,7 @@ limitations under the License. */
 
 #include <istream>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 namespace paddle {
@@ -26,19 +27,27 @@ struct OpInputConfig {
   OpInputConfig() {}
   explicit OpInputConfig(std::istream& is);
 
+  void ParseDims(std::istream& is);
+  void ParseLoD(std::istream& is);
+
   std::string name;
   std::vector<int64_t> dims;
+  std::vector<std::vector<size_t>> lod;
 };
 
 struct OpTesterConfig {
   OpTesterConfig() {}
   explicit OpTesterConfig(const std::string& filename);
-  void Init(std::istream& is);
+
+  bool Init(std::istream& is);
+
+  bool ParseAttrs(std::istream& is);
 
   const OpInputConfig* GetInput(const std::string& name);
 
   std::string op_type;
   std::vector<OpInputConfig> inputs;
+  std::unordered_map<std::string, std::string> attrs;
   int device_id{-1};  // CPU: -1
   int repeat{1};
   int profile{0};

paddle/fluid/operators/conv_op.cc

@@ -81,6 +81,7 @@ framework::OpKernelType ConvOp::GetExpectedKernelType(
       framework::OpKernelType::kDefaultCustomizedTypeValue;
   framework::LibraryType library{framework::LibraryType::kPlain};
   // TODO(pzelazko-intel): enable MKLDNN layout when it's ready
+  auto input_data_type = ctx.Input<Tensor>("Input")->type();
   std::string data_format = ctx.Attr<std::string>("data_format");
   framework::DataLayout layout = framework::StringToDataLayout(data_format);
 
@@ -94,11 +95,14 @@ framework::OpKernelType ConvOp::GetExpectedKernelType(
       platform::CanMKLDNNBeUsed(ctx)) {
     library = framework::LibraryType::kMKLDNN;
     layout = framework::DataLayout::kMKLDNN;
-    customized_type_value = kConvMKLDNNFP32;
+    customized_type_value =
+        (input_data_type == framework::DataTypeTrait<int8_t>::DataType ||
+         input_data_type == framework::DataTypeTrait<uint8_t>::DataType)
+            ? kConvMKLDNNINT8
+            : kConvMKLDNNFP32;
   }
 #endif
 
-  auto input_data_type = ctx.Input<Tensor>("Input")->type();
   if (input_data_type != framework::proto::VarType::INT8 &&
       input_data_type != framework::proto::VarType::UINT8) {
     auto filter_data_type = ctx.Input<Tensor>("Filter")->type();

paddle/fluid/operators/cross_entropy_op.cc

@@ -32,14 +32,23 @@ class CrossEntropyOp : public framework::OperatorWithKernel {
     int rank = x_dims.size();
     PADDLE_ENFORCE_EQ(rank, label_dims.size(),
                       "Input(X) and Input(Label) shall have the same rank.");
-    PADDLE_ENFORCE_EQ(framework::slice_ddim(x_dims, 0, rank - 1),
-                      framework::slice_ddim(label_dims, 0, rank - 1),
-                      "Input(X) and Input(Label) shall have the same shape "
-                      "except the last dimension.");
+    bool check = true;
+    if ((!ctx->IsRuntime()) && (framework::product(x_dims) <= 0 ||
+                                framework::product(label_dims) <= 0)) {
+      check = false;
+    }
+    if (check) {
+      PADDLE_ENFORCE_EQ(framework::slice_ddim(x_dims, 0, rank - 1),
+                        framework::slice_ddim(label_dims, 0, rank - 1),
+                        "Input(X) and Input(Label) shall have the same shape "
+                        "except the last dimension.");
+    }
     if (ctx->Attrs().Get<bool>("soft_label")) {
-      PADDLE_ENFORCE_EQ(x_dims[rank - 1], label_dims[rank - 1],
-                        "If Attr(soft_label) == true, the last dimension of "
-                        "Input(X) and Input(Label) should be equal.");
+      if (check) {
+        PADDLE_ENFORCE_EQ(x_dims[rank - 1], label_dims[rank - 1],
+                          "If Attr(soft_label) == true, the last dimension of "
+                          "Input(X) and Input(Label) should be equal.");
+      }
     } else {
       PADDLE_ENFORCE_EQ(label_dims[rank - 1], 1UL,
                         "If Attr(softLabel) == false, the last dimension of "
@@ -82,20 +91,32 @@ class CrossEntropyGradientOp : public framework::OperatorWithKernel {
                       "Input(Y@Grad) and Input(X) should have the same rank.");
     PADDLE_ENFORCE_EQ(label_dims.size(), rank,
                       "Input(Label) and Input(X) should have the same rank.");
-    PADDLE_ENFORCE_EQ(framework::slice_ddim(x_dims, 0, rank - 1),
-                      framework::slice_ddim(label_dims, 0, rank - 1),
-                      "The Input(X) and Input(Label) should have the same "
-                      "shape except the last dimension.");
-    PADDLE_ENFORCE_EQ(framework::slice_ddim(x_dims, 0, rank - 1),
-                      framework::slice_ddim(dy_dims, 0, rank - 1),
-                      "The Input(X) and Input(Y@Grad) should have the same "
-                      "shape except the last dimension.");
+
+    bool check = true;
+    if ((!ctx->IsRuntime()) && (framework::product(x_dims) <= 0 ||
+                                framework::product(label_dims) <= 0)) {
+      check = false;
+    }
+
+    if (check) {
+      PADDLE_ENFORCE_EQ(framework::slice_ddim(x_dims, 0, rank - 1),
+                        framework::slice_ddim(label_dims, 0, rank - 1),
+                        "The Input(X) and Input(Label) should have the same "
+                        "shape except the last dimension.");
+      PADDLE_ENFORCE_EQ(framework::slice_ddim(x_dims, 0, rank - 1),
+                        framework::slice_ddim(dy_dims, 0, rank - 1),
+                        "The Input(X) and Input(Y@Grad) should have the same "
+                        "shape except the last dimension.");
+    }
     PADDLE_ENFORCE_EQ(dy_dims[rank - 1], 1,
                       "The last dimension of Input(Y@Grad) should be 1.");
     if (ctx->Attrs().Get<bool>("soft_label")) {
-      PADDLE_ENFORCE_EQ(x_dims[rank - 1], label_dims[rank - 1],
-                        "When Attr(soft_label) == true, the last dimension of "
-                        "Input(X) and Input(Label) should be equal.");
+      if (check) {
+        PADDLE_ENFORCE_EQ(
+            x_dims[rank - 1], label_dims[rank - 1],
+            "When Attr(soft_label) == true, the last dimension of "
+            "Input(X) and Input(Label) should be equal.");
+      }
     } else {
       PADDLE_ENFORCE_EQ(label_dims[rank - 1], 1,
                         "When Attr(soft_label) == false, the last dimension of "

paddle/fluid/operators/data_norm_op.cc

@@ -140,9 +140,6 @@ class DataNormOpMaker : public framework::OpProtoAndCheckerMaker {
               "Scales of the history data batch, "
               "will apply to output when training")
         .AsIntermediate();
-    AddAttr<bool>("use_mkldnn",
-                  "(bool, default false) Only used in mkldnn kernel")
-        .SetDefault(false);
     AddComment(R"DOC(
 Data Normalization.
 

paddle/fluid/operators/detection/prior_box_op.h

@@ -172,6 +172,10 @@ class PriorBoxOpKernel : public framework::OpKernel<T> {
         framework::make_ddim({1, static_cast<int>(variances.size())}),
         ctx.GetPlace());
     auto var_et = framework::EigenTensor<T, 2>::From(var_t);
+
+#ifdef PADDLE_WITH_MKLML
+#pragma omp parallel for
+#endif
     for (size_t i = 0; i < variances.size(); ++i) {
       var_et(0, i) = variances[i];
     }
@@ -181,8 +185,15 @@ class PriorBoxOpKernel : public framework::OpKernel<T> {
     vars->Resize({box_num, static_cast<int>(variances.size())});
 
     auto e_vars = framework::EigenMatrix<T, Eigen::RowMajor>::From(*vars);
-    e_vars = var_et.broadcast(Eigen::DSizes<int, 2>(box_num, 1));
 
+#ifdef PADDLE_WITH_MKLML
+#pragma omp parallel for collapse(2)
+#endif
+    for (int i = 0; i < box_num; ++i) {
+      for (int j = 0; j < variances.size(); ++j) {
+        e_vars(i, j) = variances[j];
+      }
+    }
     vars->Resize(var_dim);
   }
 };  // namespace operators

paddle/fluid/operators/elementwise/mkldnn/elementwise_add_mkldnn_op.cc

@@ -77,8 +77,7 @@ class EltwiseAddMKLDNNKernel : public framework::OpKernel<T> {
       } else {
         functor.RunMidWise(n, pre, post);
       }
-      z->set_layout(DataLayout::kMKLDNN);
-      z->set_format(x->format());
+      z->set_mkldnn_prim_desc(x->get_mkldnn_prim_desc());
     } else {
       PADDLE_ENFORCE(x->layout() == DataLayout::kMKLDNN &&
                          x->format() != memory::format::format_undef,
@@ -116,7 +115,8 @@ class EltwiseAddMKLDNNKernel : public framework::OpKernel<T> {
       auto sum_pd = sum::primitive_desc(dst_md, scales, srcs_pd);
 
       // create mkldnn memory for dst
-      memory dst_memory = memory(sum_pd.dst_primitive_desc(), z_data);
+      auto dst_mem_pd = sum_pd.dst_primitive_desc();
+      memory dst_memory = memory(dst_mem_pd, z_data);
 
       std::vector<primitive::at> inputs;
       inputs.push_back(srcs[0]);
@@ -129,9 +129,7 @@ class EltwiseAddMKLDNNKernel : public framework::OpKernel<T> {
       pipeline.push_back(sum_prim);
       stream(stream::kind::eager).submit(pipeline).wait();
 
-      z->set_layout(DataLayout::kMKLDNN);
-      z->set_format(
-          (memory::format)dst_memory.get_primitive_desc().desc().data.format);
+      z->set_mkldnn_prim_desc(dst_mem_pd);
     }
   }
 };
@@ -152,24 +150,19 @@ class EltwiseAddMKLDNNGradKernel : public ElemwiseGradKernel<T> {
     auto* out = dout;
     auto *x = dout, *y = dout;
 
-    auto set_mkldnn_format = [](Tensor* in, const Tensor* out) {
-      in->set_layout(DataLayout::kMKLDNN);
-      in->set_format(out->format());
-    };
-
     if (dx != nullptr && dy != nullptr && dx->dims() == dy->dims()) {
       if (dx->dims() == dy->dims()) {
         auto blas = math::GetBlas<paddle::platform::CPUDeviceContext, T>(ctx);
         if (dx) {
           blas.VCOPY(dout->numel(), dout->data<T>(),
                      dx->mutable_data<T>(ctx.GetPlace()));
-          set_mkldnn_format(dx, dout);
+          dx->set_mkldnn_prim_desc(dout->get_mkldnn_prim_desc());
         }
 
         if (dy) {
           blas.VCOPY(dout->numel(), dout->data<T>(),
                      dy->mutable_data<T>(ctx.GetPlace()));
-          set_mkldnn_format(dy, dout);
+          dy->set_mkldnn_prim_desc(dout->get_mkldnn_prim_desc());
         }
       }
     } else {

paddle/fluid/operators/fake_quantize_op.cc

@@ -31,7 +31,7 @@ template <typename T>
 struct FindAbsMaxFunctor<platform::CPUDeviceContext, T> {
   void operator()(const platform::CPUDeviceContext& ctx, const T* in,
                   const int num, T* out) {
-    *out = *(std::max_element(in + 0, in + num, Compare<T>()));
+    *out = std::abs(*(std::max_element(in + 0, in + num, Compare<T>())));
   }
 };
 
@@ -46,10 +46,8 @@ struct ClipAndFakeQuantFunctor<platform::CPUDeviceContext, T> {
     platform::Transform<platform::CPUDeviceContext> trans;
     trans(ctx, in.data<T>(), in.data<T>() + in.numel(),
           out->mutable_data<T>(ctx.GetPlace()), ClipFunctor<T>(-s, s));
-    auto in_e = framework::EigenVector<T>::Flatten(in);
     auto out_e = framework::EigenVector<T>::Flatten(*out);
-
-    out_e.device(*ctx.eigen_device()) = (bin_cnt / s * in_e).round();
+    out_e.device(*ctx.eigen_device()) = (bin_cnt / s * out_e).round();
   }
 };
 

paddle/fluid/operators/fused/fused_embedding_seq_pool_op.cc

@@ -23,6 +23,9 @@ class FusedEmbeddingSeqPoolOp : public framework::OperatorWithKernel {
   using framework::OperatorWithKernel::OperatorWithKernel;
 
   void InferShape(framework::InferShapeContext* ctx) const override {
+    if (ctx->IsRuntime()) {
+      return;
+    }
     PADDLE_ENFORCE(ctx->HasInput("W"),
                    "Input W of FusedEmbeddingSeqPoolOp should not be null.");
     PADDLE_ENFORCE(ctx->HasInput("Ids"),
@@ -42,36 +45,15 @@ class FusedEmbeddingSeqPoolOp : public framework::OperatorWithKernel {
     // we only support sum now
     PADDLE_ENFORCE_EQ(combiner, "sum");
 
-    int64_t last_dim = table_dims[1];
-    for (int i = 1; i != ids_dims.size(); ++i) {
-      last_dim *= ids_dims[i];
-    }
-
-    if (ctx->IsRuntime()) {
-      framework::Variable* ids_var =
-          boost::get<framework::Variable*>(ctx->GetInputVarPtrs("Ids")[0]);
-      const auto& ids_lod = ids_var->Get<LoDTensor>().lod();
+    int64_t last_dim = FusedEmbeddingSeqPoolLastDim(table_dims, ids_dims);
+    // in compile time, the lod level of ids must be 1
+    framework::VarDesc* ids_desc =
+        boost::get<framework::VarDesc*>(ctx->GetInputVarPtrs("Ids")[0]);
+    PADDLE_ENFORCE_EQ(ids_desc->GetLoDLevel(), 1);
 
-      // in run time, the LoD of ids must be 1
-      PADDLE_ENFORCE(ids_lod.size(), 1u,
-                     "The LoD level of Input(Ids) must be 1");
-      PADDLE_ENFORCE_GE(ids_lod[0].size(), 1u, "The LoD could NOT be empty");
-
-      int64_t batch_size = ids_lod[0].size() - 1;
-
-      // in run time, the shape from Ids -> output
-      // should be [seq_length, 1] -> [batch_size, embedding_size]
-      ctx->SetOutputDim("Out", framework::make_ddim({batch_size, last_dim}));
-    } else {
-      // in compile time, the lod level of ids must be 1
-      framework::VarDesc* ids_desc =
-          boost::get<framework::VarDesc*>(ctx->GetInputVarPtrs("Ids")[0]);
-      PADDLE_ENFORCE_EQ(ids_desc->GetLoDLevel(), 1);
-
-      // in compile time, the shape from Ids -> output
-      // should be [-1, 1] -> [-1, embedding_size]
-      ctx->SetOutputDim("Out", framework::make_ddim({-1, last_dim}));
-    }
+    // in compile time, the shape from Ids -> output
+    // should be [-1, 1] -> [-1, embedding_size]
+    ctx->SetOutputDim("Out", framework::make_ddim({-1, last_dim}));
   }
 
  protected:

paddle/fluid/operators/fused/fused_embedding_seq_pool_op.h

@@ -61,6 +61,15 @@ struct EmbeddingVSumFunctor {
   }
 };
 
+inline int FusedEmbeddingSeqPoolLastDim(const framework::DDim &table_dims,
+                                        const framework::DDim &ids_dims) {
+  int64_t last_dim = table_dims[1];
+  for (int i = 1; i != ids_dims.size(); ++i) {
+    last_dim *= ids_dims[i];
+  }
+  return last_dim;
+}
+
 template <typename T>
 class FusedEmbeddingSeqPoolKernel : public framework::OpKernel<T> {
  public:
@@ -70,6 +79,17 @@ class FusedEmbeddingSeqPoolKernel : public framework::OpKernel<T> {
     const LoDTensor *table_var = context.Input<LoDTensor>("W");
     const std::string &combiner_type = context.Attr<std::string>("combiner");
 
+    int64_t last_dim =
+        FusedEmbeddingSeqPoolLastDim(table_var->dims(), ids_t->dims());
+    const auto &ids_lod = ids_t->lod();
+    // in run time, the LoD of ids must be 1
+    PADDLE_ENFORCE(ids_lod.size(), 1u, "The LoD level of Input(Ids) must be 1");
+    PADDLE_ENFORCE_GE(ids_lod[0].size(), 1u, "The LoD could NOT be empty");
+    int64_t batch_size = ids_lod[0].size() - 1;
+    // in run time, the shape from Ids -> output
+    // should be [seq_length, 1] -> [batch_size, embedding_size]
+    output_t->Resize({batch_size, last_dim});
+
     if (combiner_type == "sum") {
       EmbeddingVSumFunctor<T> functor;
       functor(context, table_var, ids_t, output_t);

paddle/fluid/operators/hash_op.cc

@@ -14,7 +14,6 @@ limitations under the License. */
 
 #include "paddle/fluid/operators/hash_op.h"
 #include <string>
-#include <vector>
 
 namespace paddle {
 namespace operators {
@@ -27,6 +26,9 @@ class HashOp : public framework::OperatorWithKernel {
       : OperatorWithKernel(type, inputs, outputs, attrs) {}
 
   void InferShape(framework::InferShapeContext *ctx) const override {
+    if (ctx->IsRuntime()) {
+      return;
+    }
     PADDLE_ENFORCE(ctx->HasInput("X"),
                    "Input(X) of HashOp should not be null.");
     PADDLE_ENFORCE(ctx->HasOutput("Out"),
@@ -36,15 +38,8 @@ class HashOp : public framework::OperatorWithKernel {
     PADDLE_ENFORCE_EQ(dims.size(), 2UL,
                       "The input of hash_op's dimensions must be 2");
     std::vector<int64_t> out_dims;
-    out_dims.reserve(dims.size() + 1);
-    // copy all dims except the last one
-    for (int i = 0u; i != dims.size() - 1; ++i) {
-      out_dims.emplace_back(dims[i]);
-    }
     int num_hash = ctx->Attrs().Get<int>("num_hash");
-    out_dims.emplace_back(num_hash);
-    // keep the last dim to 1
-    out_dims.emplace_back(1);
+    HashOutputSize(dims, out_dims, num_hash);
 
     ctx->SetOutputDim("Out", framework::make_ddim(out_dims));
     ctx->ShareLoD("X", /*->*/ "Out");
@@ -71,4 +66,4 @@ $$Out = scale * X$$
 namespace ops = paddle::operators;
 
 REGISTER_OP_WITHOUT_GRADIENT(hash, ops::HashOp, ops::HashOpMaker);
-REGISTER_OP_CPU_KERNEL(hash, ops::HashKerel<int>, ops::HashKerel<int64_t>);
+REGISTER_OP_CPU_KERNEL(hash, ops::HashKernel<int>, ops::HashKernel<int64_t>);

paddle/fluid/operators/hash_op.h

@@ -17,21 +17,34 @@ limitations under the License. */
 extern "C" {
 #include <xxhash.h>
 }
+#include <vector>
 #include "paddle/fluid/framework/eigen.h"
 #include "paddle/fluid/framework/op_registry.h"
 
 namespace paddle {
 namespace operators {
-// template <typename DeviceContext, typename T>
+
+inline void HashOutputSize(const framework::DDim& in_dims,
+                           std::vector<int64_t>& out_dims,  // NOLINT
+                           int num_hash) {
+  out_dims.reserve(in_dims.size() + 1);
+  // copy all dims except the last one
+  for (int i = 0u; i != in_dims.size() - 1; ++i) {
+    out_dims.emplace_back(in_dims[i]);
+  }
+  out_dims.emplace_back(num_hash);
+  // keep the last dim to 1
+  out_dims.emplace_back(1);
+}
+
 template <typename T>
-class HashKerel : public framework::OpKernel<T> {
+class HashKernel : public framework::OpKernel<T> {
  public:
   virtual void Compute(const framework::ExecutionContext& context) const {
     auto* out_t = context.Output<framework::LoDTensor>("Out");
     auto* in_t = context.Input<framework::LoDTensor>("X");
     int mod_by = context.Attr<int>("mod_by");
     int num_hash = context.Attr<int>("num_hash");
-    auto* output = out_t->mutable_data<T>(context.GetPlace());
 
     auto in_dims = in_t->dims();
     auto in_lod = in_t->lod();
@@ -39,6 +52,11 @@ class HashKerel : public framework::OpKernel<T> {
         static_cast<uint64_t>(in_dims[0]), in_lod[0].back(),
         "The actual input data's size mismatched with LoD information.");
 
+    std::vector<int64_t> out_dims;
+    HashOutputSize(in_dims, out_dims, num_hash);
+    out_t->Resize(framework::make_ddim(out_dims));
+    auto* output = out_t->mutable_data<T>(context.GetPlace());
+
     auto seq_length = in_dims[0];
     auto last_dim = in_dims[in_dims.size() - 1];
     auto* input = in_t->data<T>();
@@ -49,6 +67,7 @@ class HashKerel : public framework::OpKernel<T> {
       }
       input += last_dim;
     }
+    out_t->set_lod(in_t->lod());
   }
 };
 

paddle/fluid/operators/interpolate_op.cc

@@ -84,13 +84,13 @@ class InterpolateOpMaker : public framework::OpProtoAndCheckerMaker {
         .SetDefault("bilinear");
     AddAttr<bool>(
         "align_corners",
-        "an optinal bool. Defaults to True. "
+        "an optional bool. Defaults to True. "
         "If True, the centers of 4 corner pixels of the input and output "
         "tensors are aligned, preserving the values at the corner pixels, "
-        "if Flase, are not aligned")
+        "If False, are not aligned")
         .SetDefault(true);
     AddAttr<int>("align_mode",
-                 "(int, default \'1\'), optional for bilinear interpolation"
+                 "(int, default \'1\'), optional for bilinear interpolation, "
                  "can be \'0\' for src_idx = scale*(dst_indx+0.5)-0.5 , "
                  "can be \'1\' for src_idx = scale*dst_index .")
         .SetDefault(1);

paddle/fluid/operators/is_empty_op.cc

@@ -34,9 +34,8 @@ class IsEmptyOp : public framework::OperatorWithKernel {
 
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext &ctx) const override {
-    framework::OpKernelType kt = framework::OpKernelType(
-        ctx.Input<framework::LoDTensor>("X")->type(), platform::CPUPlace());
-    return kt;
+    auto *x = ctx.Input<framework::LoDTensor>("X");
+    return framework::OpKernelType(x->type(), x->place());
   }
 };
 
@@ -58,7 +57,6 @@ It will just return product(tensor.ddims()) > 0;
 }  // namespace paddle
 
 namespace ops = paddle::operators;
-
 REGISTER_OPERATOR(is_empty, ops::IsEmptyOp, ops::IsEmptyOpMaker,
                   paddle::framework::EmptyGradOpMaker);
 REGISTER_OP_CPU_KERNEL(

paddle/fluid/operators/is_empty_op.cu.cc

+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/operators/is_empty_op.h"
+#include "paddle/fluid/framework/op_registry.h"
+
+namespace ops = paddle::operators;
+REGISTER_OP_CUDA_KERNEL(
+    is_empty, ops::IsEmptyOpKernel<paddle::platform::CUDADeviceContext, float>,
+    ops::IsEmptyOpKernel<paddle::platform::CUDADeviceContext, double>,
+    ops::IsEmptyOpKernel<paddle::platform::CUDADeviceContext, int>,
+    ops::IsEmptyOpKernel<paddle::platform::CUDADeviceContext, int64_t>);

paddle/fluid/operators/is_empty_op.h

@@ -28,6 +28,9 @@ class IsEmptyOpKernel : public framework::OpKernel<T> {
     // get output
     auto* output_tensor = context.Output<framework::LoDTensor>("Out");
 
+    // Note: is_empty is always executed on CPU and the output data should
+    // always be allocated for CPUPlace. We reigister CUDA kernel for this op to
+    // avoid the unnecessary data transform.
     output_tensor->mutable_data<bool>(platform::CPUPlace())[0] =
         framework::product(input_tensor->dims()) == 0;
   }

paddle/fluid/operators/jit/benchmark.cc

@@ -332,6 +332,45 @@ void BenchEmbSeqPoolKernel() {
   }
 }
 
+template <jit::KernelType KT, typename T, typename PlaceType>
+void BenchSgdKernel() {
+  const T lr = 0.1;
+  auto UnDuplicatedRandomVec = [](int n, const int64_t lower,
+                                  const int64_t upper) -> std::vector<int64_t> {
+    PADDLE_ENFORCE_LE(static_cast<size_t>(upper - lower), n - 1);
+    PADDLE_ENFORCE_GT(n, 0);
+    std::vector<int64_t> all, out;
+    for (int i = 0; i < n; ++i) {
+      all.push_back(i);
+    }
+    std::random_shuffle(all.begin(), all.end());
+    out.insert(out.begin(), all.begin(), all.begin() + n);
+    return out;
+  };
+  for (int param_h : {1, 1000}) {
+    for (int grad_w : {1, 2, 8, 16, 30, 256}) {
+      // only benchmark inplace
+      Tensor param;
+      param.Resize({param_h, grad_w});
+      T* param_data = param.mutable_data<T>(PlaceType());
+      RandomVec<T>(param_h * grad_w, param_data, -2.f, 2.f);
+      for (int rows_size = 1; rows_size <= std::min(param_h, 10); ++rows_size) {
+        Tensor grad;
+        grad.Resize({rows_size, grad_w});
+        std::vector<int64_t> rows =
+            UnDuplicatedRandomVec(rows_size, 0, rows_size - 1);
+        RandomVec<T>(rows_size * grad_w, grad.mutable_data<T>(PlaceType()),
+                     -2.f, 2.f);
+        const T* grad_data = grad.data<T>();
+        const int64_t* rows_data = rows.data();
+        jit::sgd_attr_t attr(param_h, grad_w, rows_size, grad_w, rows_size);
+        BenchAllImpls<KT, jit::SgdTuples<T>, PlaceType>(
+            attr, &lr, param_data, grad_data, rows_data, param_data, &attr);
+      }
+    }
+  }
+}
+
 template <jit::KernelType KT, typename T, typename PlaceType>
 void BenchMatMulKernel() {
   for (int m : {1, 2, 3, 4}) {
@@ -477,6 +516,9 @@ BENCH_FP32_CPU(kEmbSeqPool) {
   BenchEmbSeqPoolKernel<jit::kEmbSeqPool, T, CPUPlace>();
 }
 
+// sgd function
+BENCH_FP32_CPU(kSgd) { BenchSgdKernel<jit::kSgd, T, CPUPlace>(); }
+
 // matmul
 BENCH_FP32_CPU(kMatMul) { BenchMatMulKernel<jit::kMatMul, T, CPUPlace>(); }
 

paddle/fluid/operators/jit/gen/CMakeLists.txt

@@ -32,3 +32,4 @@ USE_JITKERNEL_GEN(kSeqPool)
 USE_JITKERNEL_GEN(kHMax)
 USE_JITKERNEL_GEN(kHSum)
 USE_JITKERNEL_GEN(kEmbSeqPool)
+USE_JITKERNEL_GEN(kSgd)

paddle/fluid/operators/jit/gen/jitcode.h

@@ -31,7 +31,8 @@ namespace gen {
 // Application Binary Interface
 constexpr Xbyak::Operand::Code abi_param1(Xbyak::Operand::RDI),
     abi_param2(Xbyak::Operand::RSI), abi_param3(Xbyak::Operand::RDX),
-    abi_param4(Xbyak::Operand::RCX);
+    abi_param4(Xbyak::Operand::RCX), abi_param5(Xbyak::Operand::R8),
+    abi_param6(Xbyak::Operand::R9);
 
 constexpr Xbyak::Operand::Code g_abi_regs[] = {
     Xbyak::Operand::RBX, Xbyak::Operand::RBP, Xbyak::Operand::R12,

paddle/fluid/operators/jit/gen/sgd.cc

+/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License. */
+
+#include "paddle/fluid/operators/jit/gen/sgd.h"
+#include <stddef.h>  // offsetof
+#include <vector>
+#include "paddle/fluid/operators/jit/registry.h"
+#include "paddle/fluid/platform/cpu_info.h"
+
+namespace paddle {
+namespace operators {
+namespace jit {
+namespace gen {
+
+void SgdJitCode::genCode() {
+  preCode();
+  constexpr int block = YMM_FLOAT_BLOCK;
+  constexpr int max_num_regs = 7;
+  const int num_block = w_ / block;
+  const int num_groups = num_block / max_num_regs;
+  const size_t block_size = sizeof(float) * block;
+  const size_t width_size = w_ * sizeof(float);
+  std::vector<int> groups(num_groups, max_num_regs);
+  int rest_num_regs = num_block % max_num_regs;
+  if (rest_num_regs > 0) {
+    groups.push_back(rest_num_regs);
+  }
+
+  vbroadcastss(ymm_lr, ptr[param_lr]);
+  // protect rdx
+  mov(reg_ptr_grad_i, param_grad);
+  mov(reg_ptr_rows_i, param_rows);
+
+  mov(reg_rows_size_in_byte,
+      qword[param_attr + offsetof(sgd_attr_t, selected_rows_size)]);
+  mov(rax, sizeof(int64_t));
+  mul(reg_rows_size_in_byte);
+  mov(reg_rows_size_in_byte, rax);
+  add(reg_rows_size_in_byte, reg_ptr_rows_i);
+
+  Label l_next_row;
+  L(l_next_row);
+  {
+    mov(reg_row, qword[reg_ptr_rows_i]);
+    mov(rax, width_size);
+    mul(reg_row);
+    mov(reg_row, rax);
+
+    mov(reg_ptr_param_i, param_param);
+    mov(reg_ptr_out_i, param_out);
+    add(reg_ptr_param_i, reg_row);
+    add(reg_ptr_out_i, reg_row);
+
+    size_t w_offset = 0;
+    for (int num_regs : groups) {
+      // load grad
+      size_t inner_offfset = w_offset;
+      for (int reg_i = 0; reg_i < num_regs; ++reg_i) {
+        vmovups(ymm_t(reg_i), ptr[reg_ptr_grad_i + inner_offfset]);
+        inner_offfset += block_size;
+      }
+
+      // load param
+      inner_offfset = w_offset;
+      for (int reg_i = 0; reg_i < num_regs; ++reg_i) {
+        vmovups(ymm_t(reg_i + num_regs), ptr[reg_ptr_param_i + inner_offfset]);
+        inner_offfset += block_size;
+      }
+
+      // compute out
+      for (int reg_i = 0; reg_i < num_regs; ++reg_i) {
+        vmulps(ymm_t(reg_i), ymm_t(reg_i), ymm_lr);
+        vsubps(ymm_t(reg_i + num_regs), ymm_t(reg_i + num_regs), ymm_t(reg_i));
+      }
+
+      // save out
+      inner_offfset = w_offset;
+      for (int reg_i = 0; reg_i < num_regs; ++reg_i) {
+        vmovups(ptr[reg_ptr_out_i + inner_offfset], ymm_t(reg_i + num_regs));
+        inner_offfset += block_size;
+      }
+      w_offset += (block_size * num_regs);
+    }
+
+    add(reg_ptr_grad_i, width_size);
+    add(reg_ptr_rows_i, sizeof(int64_t));
+    cmp(reg_ptr_rows_i, reg_rows_size_in_byte);
+    jl(l_next_row, T_NEAR);
+  }
+
+  postCode();
+}
+
+class SgdCreator : public JitCodeCreator<sgd_attr_t> {
+ public:
+  bool UseMe(const sgd_attr_t& attr) const override {
+    return platform::MayIUse(platform::avx) &&
+           attr.grad_width % YMM_FLOAT_BLOCK == 0;
+  }
+  size_t CodeSize(const sgd_attr_t& attr) const override {
+    return 96 + (attr.grad_width / YMM_FLOAT_BLOCK) * 32 * 8;
+  }
+  std::unique_ptr<GenBase> CreateJitCode(
+      const sgd_attr_t& attr) const override {
+    PADDLE_ENFORCE_EQ(attr.param_width, attr.grad_width);
+    PADDLE_ENFORCE_LE(attr.selected_rows_size, attr.grad_height);
+    PADDLE_ENFORCE_GE(attr.selected_rows_size, 0);
+    return make_unique<SgdJitCode>(attr, CodeSize(attr));
+  }
+};
+
+}  // namespace gen
+}  // namespace jit
+}  // namespace operators
+}  // namespace paddle
+
+namespace gen = paddle::operators::jit::gen;
+
+REGISTER_JITKERNEL_GEN(kSgd, gen::SgdCreator);

paddle/fluid/operators/jit/gen/sgd.h

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License. */
+
+#pragma once
+
+#include <string>
+#include "glog/logging.h"
+#include "paddle/fluid/operators/jit/gen/jitcode.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace operators {
+namespace jit {
+namespace gen {
+
+class SgdJitCode : public JitCode {
+ public:
+  explicit SgdJitCode(const sgd_attr_t& attr, size_t code_size = 256 * 1024,
+                      void* code_ptr = nullptr)
+      : JitCode(code_size, code_ptr), w_(attr.grad_width) {
+    this->genCode();
+  }
+
+  DECLARE_JIT_CODE(SgdJitCode);
+  void genCode() override;
+
+ private:
+  int w_;
+  reg64_t param_lr{abi_param1};
+  reg64_t param_param{abi_param2};
+  reg64_t param_grad{abi_param3};
+  reg64_t param_rows{abi_param4};
+  reg64_t param_out{abi_param5};
+  reg64_t param_attr{abi_param6};
+
+  ymm_t ymm_lr = ymm_t(15);
+
+  reg64_t reg_ptr_grad_i{r10};
+  reg64_t reg_ptr_rows_i{r11};
+  reg64_t reg_rows_size_in_byte{r12};
+  reg64_t reg_row{r13};
+  reg64_t reg_ptr_param_i{r14};
+  reg64_t reg_ptr_out_i{r15};
+};
+
+}  // namespace gen
+}  // namespace jit
+}  // namespace operators
+}  // namespace paddle

paddle/fluid/operators/jit/helper.cc

@@ -55,6 +55,7 @@ const char* to_string(KernelType kt) {
     ONE_CASE(kHSum);
     ONE_CASE(kSoftmax);
     ONE_CASE(kEmbSeqPool);
+    ONE_CASE(kSgd);
     default:
       PADDLE_THROW("Not support type: %d, or forget to add it.", kt);
       return "NOT JITKernel";

paddle/fluid/operators/jit/helper.h

@@ -181,6 +181,14 @@ inline std::ostream& operator<<(std::ostream& os,
   return os;
 }
 
+inline std::ostream& operator<<(std::ostream& os, const sgd_attr_t& attr) {
+  os << "param_height[" << attr.param_height << "],param_width["
+     << attr.param_width << "],grad_height[" << attr.grad_height
+     << "],grad_width[" << attr.grad_width << "],selected_rows_size["
+     << attr.selected_rows_size << "]";
+  return os;
+}
+
 inline std::ostream& operator<<(std::ostream& os, const matmul_attr_t& attr) {
   os << "M[" << attr.m << "],N[" << attr.n << "],K[" << attr.k << "]";
   return os;

paddle/fluid/operators/jit/kernel_base.h

@@ -46,6 +46,7 @@ typedef enum {
   kVMul,
   kVRelu,
   kVScal,
+  kSgd,
   kVSigmoid,
   kVSquare,
   kVSub,
@@ -173,6 +174,28 @@ struct EmbSeqPoolTuples {
                             const emb_seq_pool_attr_t*);
 };
 
+typedef struct sgd_attr_s {
+  int64_t param_height, param_width;
+  int64_t grad_height, grad_width;
+  int64_t selected_rows_size;
+  sgd_attr_s() = default;
+  explicit sgd_attr_s(int64_t param_h, int64_t param_w, int64_t grad_h,
+                      int64_t grad_w, int64_t selected_rows_sz)
+      : param_height(param_h),
+        param_width(param_w),
+        grad_height(grad_h),
+        grad_width(grad_w),
+        selected_rows_size(selected_rows_sz) {}
+} sgd_attr_t;
+
+template <typename T>
+struct SgdTuples {
+  typedef T data_type;
+  typedef sgd_attr_t attr_type;
+  typedef void (*func_type)(const T*, const T*, const T*, const int64_t*, T*,
+                            const sgd_attr_t*);
+};
+
 typedef struct matmul_attr_s {
   int m, n, k;
   void* packed_weight{nullptr};

paddle/fluid/operators/jit/kernel_key.cc

@@ -13,6 +13,7 @@
  * limitations under the License. */
 
 #include "paddle/fluid/operators/jit/kernel_key.h"
+#include "paddle/fluid/platform/enforce.h"
 
 namespace paddle {
 namespace operators {
@@ -23,14 +24,30 @@ size_t JitCodeKey<int>(const int& d) {
   return d;
 }
 
+// TODO(TJ): refine and benchmark JitCodeKey generatation
 constexpr int act_type_shift = 3;  // suppot 2^3 act types
+static inline int act_type_convert(KernelType type) {
+  if (type == kVIdentity) {
+    return 0;
+  } else if (type == kVExp) {
+    return 1;
+  } else if (type == kVRelu) {
+    return 2;
+  } else if (type == kVSigmoid) {
+    return 3;
+  } else if (type == kVTanh) {
+    return 4;
+  }
+  PADDLE_THROW("Unsupported act type %d", type);
+  return 0;
+}
 
 template <>
 size_t JitCodeKey<lstm_attr_t>(const lstm_attr_t& attr) {
   size_t key = attr.d;
-  int gate_key = static_cast<int>(attr.act_gate) << 1;
-  int cand_key = static_cast<int>(attr.act_cand) << (1 + act_type_shift);
-  int cell_key = static_cast<int>(attr.act_cell) << (1 + act_type_shift * 2);
+  int gate_key = act_type_convert(attr.act_gate) << 1;
+  int cand_key = act_type_convert(attr.act_cand) << (1 + act_type_shift);
+  int cell_key = act_type_convert(attr.act_cell) << (1 + act_type_shift * 2);
   return (key << (1 + act_type_shift * 3)) + gate_key + cand_key + cell_key +
          attr.use_peephole;
 }
@@ -38,8 +55,8 @@ size_t JitCodeKey<lstm_attr_t>(const lstm_attr_t& attr) {
 template <>
 size_t JitCodeKey<gru_attr_t>(const gru_attr_t& attr) {
   size_t key = attr.d;
-  return (key << (act_type_shift * 2)) + static_cast<int>(attr.act_gate) +
-         (static_cast<int>(attr.act_cand) << act_type_shift);
+  return (key << (act_type_shift * 2)) + act_type_convert(attr.act_gate) +
+         (act_type_convert(attr.act_cand) << act_type_shift);
 }
 
 template <>
@@ -61,6 +78,11 @@ size_t JitCodeKey<emb_seq_pool_attr_t>(const emb_seq_pool_attr_t& attr) {
   return attr.table_width;
 }
 
+template <>
+size_t JitCodeKey<sgd_attr_t>(const sgd_attr_t& attr) {
+  return attr.grad_width;
+}
+
 }  // namespace jit
 }  // namespace operators
 }  // namespace paddle

paddle/fluid/operators/jit/more/mkl/CMakeLists.txt

@@ -14,3 +14,4 @@ USE_JITKERNEL_MORE(kVTanh, mkl)
 USE_JITKERNEL_MORE(kSeqPool, mkl)
 USE_JITKERNEL_MORE(kSoftmax, mkl)
 USE_JITKERNEL_MORE(kEmbSeqPool, mkl)
+USE_JITKERNEL_MORE(kSgd, mkl)

paddle/fluid/operators/jit/more/mkl/mkl.cc

@@ -184,6 +184,16 @@ bool EmbSeqPoolKernel<double>::UseMe(const emb_seq_pool_attr_t& attr) const {
   return true;
 }
 
+template <>
+bool SgdKernel<float>::UseMe(const sgd_attr_t& attr) const {
+  return true;
+}
+
+template <>
+bool SgdKernel<double>::UseMe(const sgd_attr_t& attr) const {
+  return true;
+}
+
 template <>
 bool MatMulKernel<float>::UseMe(const matmul_attr_t& attr) const {
   return platform::MayIUse(platform::avx);
@@ -239,5 +249,6 @@ REGISTER_MKL_KERNEL(kVTanh, VTanh);
 REGISTER_MKL_KERNEL(kSeqPool, SeqPool);
 REGISTER_MKL_KERNEL(kEmbSeqPool, EmbSeqPool);
 REGISTER_MKL_KERNEL(kSoftmax, Softmax);
+REGISTER_MKL_KERNEL(kSgd, Sgd);
 
 #undef REGISTER_MKL_KERNEL

paddle/fluid/operators/jit/more/mkl/mkl.h

@@ -142,6 +142,32 @@ void Softmax(const T* x, T* y, int n, int bs) {
   }
 }
 
+template <typename T>
+void Sgd(const T* lr, const T* param, const T* grad, const int64_t* rows,
+         T* out, const sgd_attr_t* attr) {
+  PADDLE_ENFORCE_EQ(attr->param_width, attr->grad_width);
+  PADDLE_ENFORCE_LE(attr->selected_rows_size, attr->grad_height);
+  T scalar = -lr[0];
+  int width = attr->grad_width;
+  if (out == param) {
+    for (int64_t i = 0; i < attr->selected_rows_size; ++i) {
+      auto h_idx = rows[i];
+      PADDLE_ENFORCE_LT(h_idx, attr->param_height);
+      PADDLE_ENFORCE_GE(h_idx, 0);
+      VAXPY(scalar, grad + i * width, out + h_idx * width, width);
+    }
+  } else {
+    for (int64_t i = 0; i < attr->selected_rows_size; ++i) {
+      auto h_idx = rows[i];
+      PADDLE_ENFORCE_LT(h_idx, attr->param_height);
+      PADDLE_ENFORCE_GE(h_idx, 0);
+      VScal(&scalar, grad + i * width, out + h_idx * width, width);
+      VAdd(param + h_idx * width, out + h_idx * width, out + h_idx * width,
+           width);
+    }
+  }
+}
+
 #define DECLARE_MKL_KERNEL(name, tuples)                             \
   template <typename T>                                              \
   class name##Kernel : public KernelMore<tuples<T>> {                \
@@ -173,6 +199,8 @@ DECLARE_MKL_KERNEL(EmbSeqPool, EmbSeqPoolTuples);
 
 DECLARE_MKL_KERNEL(Softmax, SoftmaxTuples);
 
+DECLARE_MKL_KERNEL(Sgd, SgdTuples);
+
 #undef DECLARE_MKL_KERNEL
 
 }  // namespace mkl

paddle/fluid/operators/jit/refer/CMakeLists.txt

@@ -33,3 +33,4 @@ USE_JITKERNEL_REFER(kHSum)
 USE_JITKERNEL_REFER(kHMax)
 USE_JITKERNEL_REFER(kSoftmax)
 USE_JITKERNEL_REFER(kEmbSeqPool)
+USE_JITKERNEL_REFER(kSgd)

paddle/fluid/operators/jit/refer/refer.cc

@@ -59,4 +59,6 @@ REGISTER_REFER_KERNEL(kSoftmax, Softmax);
 
 REGISTER_REFER_KERNEL(kEmbSeqPool, EmbSeqPool);
 
+REGISTER_REFER_KERNEL(kSgd, Sgd);
+
 #undef REGISTER_REFER_KERNEL

paddle/fluid/operators/jit/refer/refer.h

@@ -446,6 +446,36 @@ void EmbSeqPool(const T* table, const int64_t* idx, T* out,
   }
 }
 
+// SGD algorithm:
+// lr is pointor of learning rate scalar
+// param is an input matrix with (param_h, param_w)
+// grad is an input matrix with (grad_h, grad_w), here grad_w == param_w
+// selected_rows is a vectot<int64_t> with size selected_rows_size( <= grad_h )
+// out is an output matrix with (param_h, param_w)
+//
+// support both regular and sparse grad
+// regular SGD: out[:] = param[:] - lr[0] * grad[:];
+// sparse SGD: out[rows[i]][:] = param[rows[i]][:] - lr[0] * grad[i][:]
+//
+// Note: when use sparse SGD, and if out != param,
+// the out rows which are not selected have not beed changed, which maybe empty
+template <typename T>
+void Sgd(const T* lr, const T* param, const T* grad, const int64_t* rows,
+         T* out, const sgd_attr_t* attr) {
+  PADDLE_ENFORCE_EQ(attr->param_width, attr->grad_width);
+  PADDLE_ENFORCE_LE(attr->selected_rows_size, attr->grad_height);
+  for (int64_t i = 0; i < attr->selected_rows_size; ++i) {
+    auto h_idx = rows[i];
+    PADDLE_ENFORCE_LT(h_idx, attr->param_height);
+    PADDLE_ENFORCE_GE(h_idx, 0);
+    for (int64_t j = 0; j < attr->grad_width; ++j) {
+      out[h_idx * attr->grad_width + j] =
+          param[h_idx * attr->grad_width + j] -
+          lr[0] * grad[i * attr->grad_width + j];
+    }
+  }
+}
+
 #define DECLARE_REFER_KERNEL(name, tuples)             \
   template <typename T>                                \
   class name##Kernel : public ReferKernel<tuples<T>> { \
@@ -496,6 +526,8 @@ DECLARE_REFER_KERNEL(Softmax, SoftmaxTuples);
 
 DECLARE_REFER_KERNEL(EmbSeqPool, EmbSeqPoolTuples);
 
+DECLARE_REFER_KERNEL(Sgd, SgdTuples);
+
 #undef DECLARE_REFER_KERNEL
 
 }  // namespace refer

paddle/fluid/operators/math/beam_search.cu

@@ -119,6 +119,18 @@ __device__ __forceinline__ int SelectTopBeam(
       __syncthreads();
     }
 
+    if ((num_used_threads & 0x1) != 0) {
+      // If num_used_threads is a odd number, merge local top_beam of thread 0
+      // and num_used_threads - 1
+      if (tid_of_seq == 0) {
+        int index_in_sh = (num_used_threads - 1 + tid) * beam_size;
+        for (int i = 0; i < beam_size; i++) {
+          Insert(top_beam_local, top_beam[index_in_sh], beam_size);
+          index_in_sh++;
+        }
+      }
+    }
+
     num_used_threads = num_used_threads >> 1;
     if (tid_of_seq < num_used_threads) {
       int index_in_sh = (num_used_threads + tid) * beam_size;

paddle/fluid/operators/math/blas.h

@@ -184,6 +184,9 @@ class Blas {
   template <typename T>
   void VINV(int n, const T* a, T* y) const;
 
+  template <typename T>
+  void VMERF(int n, const T* a, T* y, int64_t mode) const;
+
  private:
   const DeviceContext& context_;
 };
@@ -290,6 +293,11 @@ class BlasT : private Blas<DeviceContext> {
     Base()->template VINV<T>(args...);
   }
 
+  template <typename... ARGS>
+  void VMERF(ARGS... args) const {
+    Base()->template VMERF<T>(args...);
+  }
+
  private:
   const Blas<DeviceContext>* Base() const {
     return static_cast<const Blas<DeviceContext>*>(this);

paddle/fluid/operators/math/blas_impl.h

@@ -123,6 +123,11 @@ struct CBlas<float> {
   static void VINV(ARGS... args) {
     platform::dynload::vsInv(args...);
   }
+
+  template <typename... ARGS>
+  static void VMERF(ARGS... args) {
+    platform::dynload::vmsErf(args...);
+  }
 };
 
 template <>
@@ -223,6 +228,11 @@ struct CBlas<double> {
   static void VINV(ARGS... args) {
     platform::dynload::vdInv(args...);
   }
+
+  template <typename... ARGS>
+  static void VMERF(ARGS... args) {
+    platform::dynload::vmdErf(args...);
+  }
 };
 
 #else
@@ -625,6 +635,19 @@ void Blas<DeviceContext>::VINV(int n, const T *a, T *y) const {
 #endif
 }
 
+template <>
+template <typename T>
+void Blas<platform::CPUDeviceContext>::VMERF(int n, const T *a, T *y,
+                                             int64_t mode) const {
+#ifdef PADDLE_WITH_MKLML
+  CBlas<T>::VMERF(n, a, y, mode);
+#else
+  for (int i = 0; i < n; ++i) {
+    y[i] = std::erf(a[i]);
+  }
+#endif
+}
+
 }  // namespace math
 }  // namespace operators
 }  // namespace paddle

paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc

@@ -96,8 +96,7 @@ void eltwise_forward(const framework::ExecutionContext &ctx,
 
   std::vector<int> src_tz = framework::vectorize2int(x->dims());
 
-  auto src_format =
-      src_tz.size() == 2 ? mkldnn::memory::format::nc : x->format();
+  auto src_format = x->format();
 
   const std::string key = gethash(src_tz, algorithm);
   const std::string key_src_data =
@@ -127,10 +126,8 @@ void eltwise_forward(const framework::ExecutionContext &ctx,
 
   if (p_fwd == nullptr) {
     // create mkldnn memory for input X
-    auto src_md = platform::MKLDNNMemDesc(
-        src_tz, platform::MKLDNNGetDataType<T>(), src_format);
     auto src_memory = std::shared_ptr<memory>(
-        new memory({src_md, mkldnn_engine}, to_void_cast(x_data)));
+        new memory(x->get_mkldnn_prim_desc(), to_void_cast(x_data)));
     // save src_memory to be referred in backward path
     dev_ctx.SetBlob(key_src_mem, src_memory);
 
@@ -177,8 +174,7 @@ void eltwise_forward(const framework::ExecutionContext &ctx,
   pipeline.push_back(*p_fwd);
   stream(stream::kind::eager).submit(pipeline).wait();
 
-  y->set_layout(DataLayout::kMKLDNN);
-  y->set_format(GetMKLDNNFormat(*dst_memory));
+  y->set_mkldnn_prim_desc(dst_memory->get_primitive_desc());
 }
 
 template <typename T>
@@ -196,9 +192,6 @@ void eltwise_grad(const framework::ExecutionContext &ctx,
 
   std::vector<int> diff_dst_tz = framework::vectorize2int(diff_y->dims());
 
-  auto diff_y_format =
-      diff_dst_tz.size() == 2 ? mkldnn::memory::format::nc : diff_y->format();
-
   const std::string key = gethash(diff_dst_tz, algorithm);
   const std::string key_src_data =
       key + ctx.op().Input("Out") + "@eltwise_fwd_src_data";
@@ -210,8 +203,8 @@ void eltwise_grad(const framework::ExecutionContext &ctx,
       key + std::to_string(*p_src_layout) + "@eltwise_fwd_src_mem";
   const std::string key_fwd_pd =
       key + std::to_string(*p_src_layout) + "@eltwise_fwd_pd";
-  const std::string key_with_layouts =
-      key + std::to_string(*p_src_layout) + "-" + std::to_string(diff_y_format);
+  const std::string key_with_layouts = key + std::to_string(*p_src_layout) +
+                                       "-" + std::to_string(diff_y->format());
   const std::string key_diff_src_mem =
       key_with_layouts + "@eltwise_diff_src_mem";
   const std::string key_diff_dst_mem =
@@ -234,10 +227,8 @@ void eltwise_grad(const framework::ExecutionContext &ctx,
 
   if (p_grad == nullptr) {
     // create mkldnn memory for input diff_y
-    auto diff_dst_md = platform::MKLDNNMemDesc(
-        diff_dst_tz, platform::MKLDNNGetDataType<T>(), diff_y_format);
     auto diff_dst_memory = std::shared_ptr<memory>(
-        new memory({diff_dst_md, mkldnn_engine}, to_void_cast(diff_y_data)));
+        new memory(diff_y->get_mkldnn_prim_desc(), to_void_cast(diff_y_data)));
     dev_ctx.SetBlob(key_diff_dst_mem, diff_dst_memory);
 
     // retrieve eltwise primitive desc from device context
@@ -281,8 +272,7 @@ void eltwise_grad(const framework::ExecutionContext &ctx,
   pipeline.push_back(*p_grad);
   stream(stream::kind::eager).submit(pipeline).wait();
 
-  diff_x->set_layout(DataLayout::kMKLDNN);
-  diff_x->set_format(GetMKLDNNFormat(*diff_src_memory));
+  diff_x->set_mkldnn_prim_desc(diff_src_memory->get_primitive_desc());
 }
 
 template <typename T, mkldnn::algorithm algorithm>

paddle/fluid/operators/mkldnn/batch_norm_mkldnn_op.cc

@@ -206,17 +206,14 @@ class BatchNormMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
     if (fuse_with_relu) flags |= mkldnn::fuse_bn_relu;
 
     // create mkldnn memory from input x tensor
-    mkldnn::memory::format input_format =
-        platform::MKLDNNFormatForSize(src_tz.size(), x->format());
 
     // keys for backward pass
     const std::string key = BatchNormMKLDNNHandler::GetHash(
-        src_tz, epsilon, flags, global_stats, input_format,
+        src_tz, epsilon, flags, global_stats, x->format(),
         ctx.op().Output("SavedMean"));
     const std::string key_batch_norm_fwd_pd = key + "@bn_fwd_pd";
 
-    auto user_src_md = platform::MKLDNNMemDesc(
-        {src_tz}, platform::MKLDNNGetDataType<T>(), input_format);
+    auto user_src_md = x->get_mkldnn_prim_desc().desc();
 
     // create primitive descriptor for batch norm forward
     using bn_fwd_types = bn_type_traits<mkldnn::batch_normalization_forward>;
@@ -230,8 +227,8 @@ class BatchNormMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
     BatchNormMKLDNNHandler handler(batch_norm_fwd_pd, dev_ctx, mkldnn_engine,
                                    key);
 
-    auto src_memory =
-        handler.AcquireSrcMemory(user_src_md, to_void_cast(x_data));
+    auto src_memory = handler.AcquireSrcMemory(x->get_mkldnn_prim_desc(),
+                                               to_void_cast(x_data));
 
     // crate mkldnn memory for weights(scale/shift)
     auto scaleshift_memory =
@@ -265,8 +262,7 @@ class BatchNormMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
           variance_memory, false);
     }
 
-    y->set_layout(DataLayout::kMKLDNN);
-    y->set_format(platform::GetMKLDNNFormat(*dst_memory));
+    y->set_mkldnn_prim_desc(dst_memory->get_primitive_desc());
 
     std::vector<mkldnn::primitive> pipeline;
     pipeline.push_back(*batch_norm_p);
@@ -336,9 +332,6 @@ class BatchNormMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
 
     using bn_bwd_types = bn_type_traits<mkldnn::batch_normalization_backward>;
 
-    mkldnn::memory::format dst_format =
-        platform::MKLDNNFormatForSize(src_tz.size(), diff_y->format());
-
     mkldnn::memory::format input_format =
         platform::MKLDNNFormatForSize(src_tz.size(), x->format());
 
@@ -346,14 +339,14 @@ class BatchNormMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
 
     // keys from forward pass
     const std::string key = BatchNormMKLDNNHandler::GetHash(
-        src_tz, epsilon, flags, false, input_format,
+        src_tz, epsilon, flags, false, x->format(),
         ctx.op().Input("SavedMean"));
     const std::string key_batch_norm_fwd_pd = key + "@bn_fwd_pd";
 
     // keys for primitives reuse
     const std::string key_with_hash =
         key + BatchNormMKLDNNHandler::GetHash(src_tz, epsilon, flags, false,
-                                              input_format);
+                                              x->format());
     const std::string key_batch_norm_bwd_p =
         key_with_hash + "@batch_norm_bwd_p";
     const std::string key_batch_norm_src_mem_p =
@@ -373,9 +366,8 @@ class BatchNormMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
 
     primitive reorder_diff_dst;
     bool is_diff_dst_reordered = false;
-    auto user_diff_dst_memory = memory(
-        {{{diff_dst_tz}, memory::data_type::f32, dst_format}, mkldnn_engine},
-        to_void_cast(diff_y_data));
+    auto user_diff_dst_memory =
+        memory(diff_y->get_mkldnn_prim_desc(), to_void_cast(diff_y_data));
 
     // MKLDNN requires a single piece of memory for scale and shift/bias data
     const size_t scaleshift_size = 2 * ic;
@@ -459,10 +451,7 @@ class BatchNormMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
       dev_ctx.SetBlob(key_batch_norm_diff_dst_mem_p, diff_dst_memory);
 
       // set layout/format of output tensors
-      diff_x->set_layout(DataLayout::kMKLDNN);
-      diff_x->set_format((memory::format)diff_src_memory->get_primitive_desc()
-                             .desc()
-                             .data.format);
+      diff_x->set_mkldnn_prim_desc(diff_src_memory->get_primitive_desc());
     } else {
       // primitives already exist
       UpdateMemoryData(dev_ctx, key_batch_norm_src_mem_p, to_void_cast(x_data));
@@ -487,10 +476,7 @@ class BatchNormMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
       }
 
       // set layout/format of output tensors
-      diff_x->set_layout(DataLayout::kMKLDNN);
-      diff_x->set_format((memory::format)diff_src_memory->get_primitive_desc()
-                             .desc()
-                             .data.format);
+      diff_x->set_mkldnn_prim_desc(diff_src_memory->get_primitive_desc());
     }
 
     // execute optional reorder and batch_norm backward primitive

paddle/fluid/operators/mkldnn/concat_mkldnn_op.cc

@@ -47,11 +47,6 @@ static memory::primitive_desc CreateMemPrimDesc(const Tensor& input,
   return mem_prim_desc;
 }
 
-static mkldnn::memory::format GetDstMemFormat(
-    const concat::primitive_desc& concat_pd) {
-  return (memory::format)concat_pd.dst_primitive_desc().desc().data.format;
-}
-
 static platform::CPUPlace GetCpuPlace(
     const paddle::framework::ExecutionContext& ctx) {
   auto place = ctx.GetPlace();
@@ -139,8 +134,7 @@ class ConcatMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
     auto concat = prim_creator.CreateConcatPrimitive(concat_pd, output, place);
     stream(stream::kind::eager).submit({concat}).wait();
 
-    output->set_layout(DataLayout::kMKLDNN);
-    output->set_format(GetDstMemFormat(concat_pd));
+    output->set_mkldnn_prim_desc(concat_pd.dst_primitive_desc());
   }
 };
 }  // namespace operators

paddle/fluid/operators/mkldnn/conv_mkldnn_op.cc

@@ -282,8 +282,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
     pipeline.push_back(*conv_p);
     stream(stream::kind::eager).submit(pipeline).wait();
 
-    auto dst_mpd = dst_memory_p->get_primitive_desc();
-    output->set_mkldnn_prim_desc(dst_mpd);
+    output->set_mkldnn_prim_desc(dst_memory_p->get_primitive_desc());
   }
   void ComputeINT8(const paddle::framework::ExecutionContext& ctx) const {
     const bool is_test = ctx.Attr<bool>("is_test");
@@ -972,8 +971,7 @@ class ConvMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
 
       pipeline.push_back(*conv_bwd_data_p);
 
-      input_grad->set_layout(DataLayout::kMKLDNN);
-      input_grad->set_format(GetMKLDNNFormat(*diff_src_memory_p));
+      input_grad->set_mkldnn_prim_desc(diff_src_memory_p->get_primitive_desc());
     }
     stream(stream::kind::eager).submit(pipeline).wait();
   }
@@ -991,12 +989,12 @@ REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv2d, MKLDNN,
 
 REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv2d, MKLDNN,
                                     ::paddle::platform::CPUPlace, U8,
-                                    ops::kConvMKLDNNFP32,
+                                    ops::kConvMKLDNNINT8,
                                     ops::ConvMKLDNNOpKernel<uint8_t, float>);
 
 REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv2d, MKLDNN,
                                     ::paddle::platform::CPUPlace, S8,
-                                    ops::kConvMKLDNNFP32,
+                                    ops::kConvMKLDNNINT8,
                                     ops::ConvMKLDNNOpKernel<int8_t, float>);
 
 REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(conv2d_grad, MKLDNN,

paddle/fluid/operators/mkldnn/conv_transpose_mkldnn_op.cc

@@ -221,8 +221,7 @@ class ConvTransposeMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
     pipeline.push_back(*conv_p);
     mkldnn::stream(mkldnn::stream::kind::eager).submit(pipeline).wait();
 
-    output->set_layout(DataLayout::kMKLDNN);
-    output->set_format(platform::GetMKLDNNFormat(*dst_memory_p));
+    output->set_mkldnn_prim_desc(dst_memory_p->get_primitive_desc());
   }
 
  private:

paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc

@@ -158,6 +158,14 @@ class SoftmaxMKLDNNKernel : public paddle::framework::OpKernel<T> {
     auto softmax_p =
         handler.AcquireSoftmax(softmax_dst_memory_p, softmax_src_memory_p);
 
+    // We cannot use softmax_dst_memory_p to get prim desc as
+    // it contains flattened dims (2D) while output tensor can
+    // have 2,3,4+ dims
+    auto output_mem_pd = paddle::platform::create_prim_desc_from_dims(
+        paddle::framework::vectorize2int(output->dims()),
+        mkldnn::memory::format::blocked);
+    output->set_mkldnn_prim_desc(output_mem_pd);
+
     std::vector<primitive> pipeline{
         *(static_cast<softmax_forward::primitive*>(softmax_p.get()))};
     stream(stream::kind::eager).submit(pipeline).wait();

paddle/fluid/operators/ngraph/ngraph_bridge.cc

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #include <algorithm>
 #include <functional>
+#include <memory>
 #include <vector>
 
 #include "ngraph/ngraph.hpp"