Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into async_executor

.gitignore

@@ -4,6 +4,7 @@ paddle/operators/tensor.save
 python/paddle/v2/fluid/tests/book/image_classification_resnet.inference.model/
 python/paddle/v2/fluid/tests/book/image_classification_vgg.inference.model/
 python/paddle/v2/fluid/tests/book/label_semantic_roles.inference.model/
+paddle/fluid/operators/distributed/send_recv.proto
 *.DS_Store
 *.vs
 build/
@@ -28,4 +29,5 @@ third_party/
 build_*
 # clion workspace.
 cmake-build-*
+paddle/fluid/operators/distributed/send_recv.proto
 model_test

CMakeLists.txt

@@ -302,6 +302,14 @@ set(PADDLE_PYTHON_BUILD_DIR "${CMAKE_CURRENT_BINARY_DIR}/python/build")
 set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "-O3 -g -DNDEBUG")
 set(CMAKE_C_FLAGS_RELWITHDEBINFO "-O3 -g -DNDEBUG")
 
+if (ON_INFER)
+    message(STATUS "On inference mode, will take place some specific optimization.")
+    add_definitions(-DPADDLE_ON_INFERENCE)
+else()
+    #TODO(luotao), combine this warning with `make inference_lib_dist` command.
+    message(WARNING "On inference mode, will take place some specific optimization. Turn on the ON_INFER flag when building inference_lib only.")
+endif()
+
 add_subdirectory(paddle)
 if(WITH_PYTHON)
     add_subdirectory(python)
@@ -312,10 +320,3 @@ if(WITH_DOC)
     find_python_module(recommonmark REQUIRED)
     add_subdirectory(doc)
 endif()
-
-if (ON_INFER)
-    message(STATUS "On inference mode, will take place some specific optimization.")
-else()
-    #TODO(luotao), combine this warning with `make inference_lib_dist` command.
-    message(WARNING "On inference mode, will take place some specific optimization. Turn on the ON_INFER flag when building inference_lib only.")
-endif()

cmake/operators.cmake

@@ -111,7 +111,7 @@ function(op_library TARGET)
 
     # Define operators that don't need pybind here.
     foreach(manual_pybind_op "compare_op" "logical_op" "nccl_op"
-"tensor_array_read_write_op" "tensorrt_engine_op")
+"tensor_array_read_write_op" "tensorrt_engine_op" "conv_fusion_op")
         if ("${TARGET}" STREQUAL "${manual_pybind_op}")
             set(pybind_flag 1)
         endif()

paddle/fluid/framework/details/exception_holder.h

@@ -30,6 +30,8 @@ class ExceptionHolder {
       Catch(exp);
     } catch (platform::EnforceNotMet exp) {
       Catch(exp);
+    } catch (std::exception& ex) {
+      LOG(FATAL) << "std::exception caught, " << ex.what();
     } catch (...) {
       LOG(FATAL) << "Unknown exception caught";
     }

paddle/fluid/framework/executor.cc

@@ -418,11 +418,6 @@ void Executor::RunPreparedContext(ExecutorPrepareContext* ctx, Scope* scope,
       DeleteUnusedTensors(*local_scope, op.get(), gc.get(),
                           &(ctx->cur_ref_cnts_));
     }
-
-    if (FLAGS_benchmark) {
-      VLOG(20) << "Memory used after operator " + op->Type() + " running: "
-               << memory::memory_usage(place_);
-    }
   }
 
   if (gc != nullptr) {
@@ -444,13 +439,6 @@ void Executor::RunPreparedContext(ExecutorPrepareContext* ctx, Scope* scope,
       scope->DropKids();
     }
   }
-
-  if (FLAGS_benchmark) {
-    VLOG(20) << "-------------------------------------------------------";
-    VLOG(20) << "Memory used after deleting local scope: "
-             << memory::memory_usage(place_);
-    VLOG(20) << "-------------------------------------------------------";
-  }
 }
 
 void Executor::RunPreparedContext(

paddle/fluid/framework/ir/conv_elementwise_add_mkldnn_fuse_pass.cc

@@ -14,14 +14,15 @@
 
 #include "paddle/fluid/framework/ir/conv_elementwise_add_mkldnn_fuse_pass.h"
 #include <functional>
-#include <utility>
+#include <list>
+#include <map>
+#include <tuple>
 
 #include "paddle/fluid/framework/ir/graph_traits.h"
 
 namespace paddle {
 namespace framework {
 namespace ir {
-namespace {
 
 // The function keeps the graph consistent by replacing
 // a node 'from' in the set of inputs nodes
@@ -51,99 +52,382 @@ void CorrectGraphEdges(Graph* graph, Node* from, Node* to) {
     }
   }
 }
-}  // namespace
-using graph_ptr = std::unique_ptr<ir::Graph>;
 
-graph_ptr ConvElementwiseAddMKLDNNFusePass::ApplyImpl(graph_ptr graph) const {
-  FusePassBase::Init(name_scope_, graph.get());
+bool IsReachable(ir::Graph* graph, Node* from, Node* to) {
+  auto find_node = [](ir::Graph* graph, const Node* node) -> Node* {
+    for (auto n : graph->Nodes()) {
+      if (n == node) {
+        return n;
+      }
+    }
 
-  GraphPatternDetector gpd;
-  auto pattern = gpd.mutable_pattern();
+    return nullptr;
+  };
 
-  patterns::Conv conv_pattern{pattern, name_scope_};
-  auto conv_output = conv_pattern();
+  if (from == to) {
+    return true;
+  }
 
-  patterns::ElementwiseAdd elementwise_add_pattern{pattern, name_scope_};
-  elementwise_add_pattern(conv_output);
+  std::map<Node*, bool> visited;
 
-  conv_output->AsIntermediate();
+  for (auto& node : GraphTraits::DFS(*graph)) {
+    visited[&node] = false;
+  }
 
-  auto conv_op_has_bias = [](const Node& conv_op) -> std::pair<bool, Node*> {
-    auto bias_input_names = conv_op.Op()->Inputs();
-    auto bias_it = bias_input_names.find("Bias");
-
-    if (bias_it != std::end(bias_input_names)) {
-      bool has_bias = !bias_it->second.empty();
-
-      if (has_bias) {
-        auto conv_bias_names = bias_it->second;
-        auto conv_bias_names_it =
-            std::find_if(std::begin(conv_op.inputs), std::end(conv_op.inputs),
-                         [&conv_bias_names](Node* n) -> bool {
-                           return n->Name() == conv_bias_names[0];
-                         });
-        return std::make_pair(has_bias, *conv_bias_names_it);
-      }
-    }
+  visited[from] = true;
 
-    return std::make_pair(false, nullptr);
-  };
+  std::list<Node*> queue;
+  queue.push_back(from);
 
-  auto handler = [&](const GraphPatternDetector::subgraph_t& subgraph,
-                     Graph* g) {
-    GET_IR_NODE_FROM_SUBGRAPH(conv_op, conv_op, conv_pattern);
-    GET_IR_NODE_FROM_SUBGRAPH(conv_input, conv_input, conv_pattern);
-    GET_IR_NODE_FROM_SUBGRAPH(conv_filter, conv_filter, conv_pattern);
-    GET_IR_NODE_FROM_SUBGRAPH(conv_output, conv_output, conv_pattern);
-    GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_op, elementwise_add_op,
-                              elementwise_add_pattern);
-    GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_x, elementwise_add_x,
-                              elementwise_add_pattern);
-    GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_out, elementwise_add_out,
-                              elementwise_add_pattern);
+  while (!queue.empty()) {
+    auto cur = find_node(graph, queue.front());
+    queue.pop_front();
 
-    if (FindFuseOption(*conv_op, *elementwise_add_op) != FUSE_MKLDNN) return;
+    if (!cur) return false;
 
-    OpDesc op_desc;
-    op_desc.SetType("conv2d");
+    for (auto n : cur->outputs) {
+      if (n == to) {
+        return true;
+      }
 
-    op_desc.SetInput("Input", {conv_input->Name()});
-    op_desc.SetInput("Filter", {conv_filter->Name()});
-    op_desc.SetInput("ResidualData", {elementwise_add_x->Name()});
-    op_desc.SetOutput("Output", {conv_output->Name()});
+      if (!visited[n]) {
+        visited[n] = true;
+        queue.push_back(n);
+      }
+    }
+  }
+  return false;
+}
 
-    bool has_bias;
-    Node* conv_bias;
+boost::optional<Node*> HasBias(const Node& op, const std::string& bias_name) {
+  auto bias_input_names = op.Op()->Inputs();
+  auto bias_it = bias_input_names.find(bias_name);
 
-    std::tie(has_bias, conv_bias) = conv_op_has_bias(*conv_op);
+  if (bias_it != std::end(bias_input_names)) {
+    bool has_bias = !bias_it->second.empty();
 
     if (has_bias) {
-      op_desc.SetInput("Bias", {conv_bias->Name()});
+      auto bias_names = bias_it->second;
+      auto bias_names_it =
+          std::find_if(std::begin(op.inputs), std::end(op.inputs),
+                       [&bias_names](Node* n) -> bool {
+                         return n->Name() == bias_names[0];
+                       });
+      return *bias_names_it;
     }
+  }
 
-    for (const auto& attr : conv_op->Op()->GetAttrMap()) {
-      op_desc.SetAttr(attr.first, attr.second);
-    }
+  return boost::none;
+}
 
-    op_desc.SetAttr("fuse_residual_connection", true);
+ResidualConnectionMKLDNNFusePass::IdentityFuseHandle::IdentityFuseHandle(
+    const ResidualConnectionMKLDNNFusePass::CanFuseFunc& can_fuse_func,
+    const ResidualConnectionMKLDNNFusePass::IdentityConvFunc&
+        get_node_from_conv_op,
+    const ResidualConnectionMKLDNNFusePass::IdentityElementwiseAddFunc&
+        get_node_from_elementwise_add_op)
+    : fusion_stats{std::make_shared<int>(0)},
+      can_fuse_func{can_fuse_func},
+      get_node_from_conv_op{get_node_from_conv_op},
+      get_node_from_elementwise_add_op{get_node_from_elementwise_add_op} {}
+
+void ResidualConnectionMKLDNNFusePass::IdentityFuseHandle::operator()(
+    const GraphPatternDetector::subgraph_t& subgraph, Graph* graph) {
+  Node* conv_op;
+  Node* conv_input;
+  Node* conv_filter;
+  Node* conv_output;
+
+  Node* elementwise_add_op;
+  Node* elementwise_add_identity;
+  Node* elementwise_add_out;
+
+  std::tie(conv_op, conv_input, conv_filter, conv_output) =
+      get_node_from_conv_op(subgraph);
+  std::tie(elementwise_add_op, elementwise_add_identity, elementwise_add_out) =
+      get_node_from_elementwise_add_op(subgraph);
+
+  if (!can_fuse_func(conv_op, elementwise_add_op)) return;
+
+  if (!IsReachable(graph, elementwise_add_identity, conv_output)) return;
+
+  OpDesc op_desc;
+  op_desc.SetType("conv2d");
+
+  op_desc.SetInput("Input", {conv_input->Name()});
+  op_desc.SetInput("Filter", {conv_filter->Name()});
+  op_desc.SetInput("ResidualData", {elementwise_add_identity->Name()});
+  op_desc.SetOutput("Output", {conv_output->Name()});
+
+  auto conv_bias = HasBias(*conv_op, "Bias");
+
+  if (conv_bias) {
+    op_desc.SetInput("Bias", {(*conv_bias)->Name()});
+  }
 
-    auto fused_conv_op = g->CreateOpNode(&op_desc);
+  for (const auto& attr : conv_op->Op()->GetAttrMap()) {
+    op_desc.SetAttr(attr.first, attr.second);
+  }
 
-    IR_NODE_LINK_TO(conv_input, fused_conv_op);
-    IR_NODE_LINK_TO(conv_filter, fused_conv_op);
-    IR_NODE_LINK_TO(elementwise_add_x, fused_conv_op);
-    IR_NODE_LINK_TO(fused_conv_op, conv_output);
+  op_desc.SetAttr("fuse_residual_connection", true);
 
-    if (has_bias) {
-      IR_NODE_LINK_TO(conv_bias, fused_conv_op);
-    }
+  auto fused_conv_op = graph->CreateOpNode(&op_desc);
 
-    CorrectGraphEdges(g, elementwise_add_out, conv_output);
-    GraphSafeRemoveNodes(g, {elementwise_add_out, conv_op, elementwise_add_op});
-  };
+  IR_NODE_LINK_TO(conv_input, fused_conv_op);
+  IR_NODE_LINK_TO(conv_filter, fused_conv_op);
+  IR_NODE_LINK_TO(elementwise_add_identity, fused_conv_op);
+  IR_NODE_LINK_TO(fused_conv_op, conv_output);
 
-  gpd(graph.get(), handler);
+  if (conv_bias) {
+    IR_NODE_LINK_TO((*conv_bias), fused_conv_op);
+  }
 
+  CorrectGraphEdges(graph, elementwise_add_out, conv_output);
+  GraphSafeRemoveNodes(graph,
+                       {elementwise_add_out, conv_op, elementwise_add_op});
+  (*fusion_stats)++;
+}
+
+ResidualConnectionMKLDNNFusePass::ProjectionFuseHandle::ProjectionFuseHandle(
+    const ResidualConnectionMKLDNNFusePass::CanFuseFunc& can_fuse_func,
+    const ResidualConnectionMKLDNNFusePass::ProjectionConvFunc&
+        get_node_from_conv_x_op,
+    const ResidualConnectionMKLDNNFusePass::ProjectionConvFunc&
+        get_node_from_conv_y_op,
+    const ResidualConnectionMKLDNNFusePass::ProjectionElementwiseAddFunc&
+        get_node_from_elementwise_add_op)
+    : fusion_stats{std::make_shared<int>(0)},
+      can_fuse_func{can_fuse_func},
+      get_node_from_conv_x_op{get_node_from_conv_x_op},
+      get_node_from_conv_y_op{get_node_from_conv_y_op},
+      get_node_from_elementwise_add_op{get_node_from_elementwise_add_op} {}
+
+void ResidualConnectionMKLDNNFusePass::ProjectionFuseHandle::operator()(
+    const GraphPatternDetector::subgraph_t& subgraph, Graph* graph) {
+  Node* conv_x_op;
+  Node* conv_x_input;
+  Node* conv_x_filter;
+  Node* conv_x_output;
+
+  Node* conv_y_op;
+  Node* conv_y_input;
+  Node* conv_y_filter;
+  Node* conv_y_output;
+
+  Node* elementwise_add_op;
+  Node* elementwise_add_out;
+
+  std::tie(conv_x_op, conv_x_input, conv_x_filter, conv_x_output) =
+      get_node_from_conv_x_op(subgraph);
+  std::tie(conv_y_op, conv_y_input, conv_y_filter, conv_y_output) =
+      get_node_from_conv_y_op(subgraph);
+  std::tie(elementwise_add_op, elementwise_add_out) =
+      get_node_from_elementwise_add_op(subgraph);
+
+  if (!can_fuse_func(conv_x_op, elementwise_add_op)) return;
+  if (!can_fuse_func(conv_y_op, elementwise_add_op)) return;
+
+  Node* projection_node;
+  Node* residual_conv_op;
+  Node* residual_conv_input;
+  Node* residual_conv_filter;
+  Node* residual_conv_output;
+
+  if (IsReachable(graph, conv_x_input, conv_y_output)) {
+    projection_node = conv_x_output;
+    residual_conv_op = conv_y_op;
+    residual_conv_input = conv_y_input;
+    residual_conv_filter = conv_y_filter;
+    residual_conv_output = conv_y_output;
+  } else if (IsReachable(graph, conv_y_input, conv_x_output)) {
+    projection_node = conv_y_output;
+    residual_conv_op = conv_x_op;
+    residual_conv_input = conv_x_input;
+    residual_conv_filter = conv_x_filter;
+    residual_conv_output = conv_x_output;
+  } else {
+    return;
+  }
+
+  OpDesc op_desc;
+  op_desc.SetType("conv2d");
+
+  op_desc.SetInput("Input", {residual_conv_input->Name()});
+  op_desc.SetInput("Filter", {residual_conv_filter->Name()});
+  op_desc.SetInput("ResidualData", {projection_node->Name()});
+  op_desc.SetOutput("Output", {residual_conv_output->Name()});
+
+  auto residual_conv_bias = HasBias(*residual_conv_op, "Bias");
+
+  if (residual_conv_bias) {
+    op_desc.SetInput("Bias", {(*residual_conv_bias)->Name()});
+  }
+
+  for (const auto& attr : residual_conv_op->Op()->GetAttrMap()) {
+    op_desc.SetAttr(attr.first, attr.second);
+  }
+
+  op_desc.SetAttr("fuse_residual_connection", true);
+
+  auto fused_conv_op = graph->CreateOpNode(&op_desc);
+
+  IR_NODE_LINK_TO(residual_conv_input, fused_conv_op);
+  IR_NODE_LINK_TO(residual_conv_filter, fused_conv_op);
+  IR_NODE_LINK_TO(projection_node, fused_conv_op);
+  IR_NODE_LINK_TO(fused_conv_op, residual_conv_output);
+
+  if (residual_conv_bias) {
+    IR_NODE_LINK_TO((*residual_conv_bias), fused_conv_op);
+  }
+
+  CorrectGraphEdges(graph, elementwise_add_out, residual_conv_output);
+  GraphSafeRemoveNodes(
+      graph, {elementwise_add_out, residual_conv_op, elementwise_add_op});
+  (*fusion_stats)++;
+}
+
+std::tuple<Node*, Node*, Node*, Node*>
+ResidualConnectionMKLDNNFusePass::GetNodesFromConv(
+    const patterns::Conv& conv_pattern,
+    const GraphPatternDetector::subgraph_t& subgraph) const {
+  GET_IR_NODE_FROM_SUBGRAPH(conv_op, conv_op, conv_pattern);
+  GET_IR_NODE_FROM_SUBGRAPH(conv_input, conv_input, conv_pattern);
+  GET_IR_NODE_FROM_SUBGRAPH(conv_filter, conv_filter, conv_pattern);
+  GET_IR_NODE_FROM_SUBGRAPH(conv_output, conv_output, conv_pattern);
+
+  return std::make_tuple(conv_op, conv_input, conv_filter, conv_output);
+}
+
+GraphWithStats ResidualConnectionMKLDNNFusePass::FuseConvAsX(
+    const std::string& name_scope,
+    const GraphWithStats& graph_with_stats) const {
+  ir::Graph* graph;
+  int stats;
+
+  std::tie(graph, stats) = graph_with_stats;
+
+  GraphPatternDetector gpd;
+  auto pattern = gpd.mutable_pattern();
+
+  patterns::Conv conv_pattern{pattern, name_scope};
+  auto conv_output = conv_pattern();
+
+  patterns::ElementwiseAdd elementwise_add_pattern{pattern, name_scope};
+  elementwise_add_pattern(
+      conv_output,
+      pattern->NewNode(elementwise_add_pattern.elementwise_add_y_repr()));
+  conv_output->AsIntermediate();
+
+  auto get_node_from_elementwise_add = [&elementwise_add_pattern](
+      const GraphPatternDetector::subgraph_t& subgraph)
+      -> std::tuple<Node*, Node*, Node*> {
+        GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_op, elementwise_add_op,
+                                  elementwise_add_pattern);
+        GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_y, elementwise_add_y,
+                                  elementwise_add_pattern);
+        GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_out, elementwise_add_out,
+                                  elementwise_add_pattern);
+
+        return std::make_tuple(elementwise_add_op, elementwise_add_y,
+                               elementwise_add_out);
+      };
+
+  return ExecuteHandleOnGraph<IdentityFuseHandle>(
+      &gpd, graph_with_stats,
+      [this, &conv_pattern](const GraphPatternDetector::subgraph_t& subgraph) {
+        return GetNodesFromConv(conv_pattern, subgraph);
+      },
+      get_node_from_elementwise_add);
+}
+
+GraphWithStats ResidualConnectionMKLDNNFusePass::FuseConvAsY(
+    const std::string& name_scope,
+    const GraphWithStats& graph_with_stats) const {
+  GraphPatternDetector gpd;
+  auto pattern = gpd.mutable_pattern();
+
+  patterns::Conv conv_pattern{pattern, name_scope};
+  auto conv_output = conv_pattern();
+
+  patterns::ElementwiseAdd elementwise_add_pattern{pattern, name_scope};
+  elementwise_add_pattern(
+      pattern->NewNode(elementwise_add_pattern.elementwise_add_x_repr()),
+      conv_output);
+  conv_output->AsIntermediate();
+
+  auto get_node_from_elementwise_add = [&elementwise_add_pattern](
+      const GraphPatternDetector::subgraph_t& subgraph)
+      -> std::tuple<Node*, Node*, Node*> {
+        GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_op, elementwise_add_op,
+                                  elementwise_add_pattern);
+        GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_x, elementwise_add_x,
+                                  elementwise_add_pattern);
+        GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_out, elementwise_add_out,
+                                  elementwise_add_pattern);
+
+        return std::make_tuple(elementwise_add_op, elementwise_add_x,
+                               elementwise_add_out);
+      };
+
+  return ExecuteHandleOnGraph<IdentityFuseHandle>(
+      &gpd, graph_with_stats,
+      [this, &conv_pattern](const GraphPatternDetector::subgraph_t& subgraph) {
+        return GetNodesFromConv(conv_pattern, subgraph);
+      },
+      get_node_from_elementwise_add);
+}
+
+GraphWithStats ResidualConnectionMKLDNNFusePass::FuseProjectionConv(
+    const std::string& name_scope,
+    const GraphWithStats& graph_with_stats) const {
+  GraphPatternDetector gpd;
+  auto pattern = gpd.mutable_pattern();
+
+  patterns::Conv conv_x_pattern{pattern, name_scope};
+  auto conv_x_output = conv_x_pattern();
+
+  patterns::Conv conv_y_pattern{pattern, name_scope};
+  auto conv_y_output = conv_y_pattern();
+
+  patterns::ElementwiseAdd elementwise_add_pattern{pattern, name_scope};
+  elementwise_add_pattern(conv_x_output, conv_y_output);
+  conv_x_output->AsIntermediate();
+  conv_y_output->AsIntermediate();
+
+  auto get_node_from_elementwise_add = [&elementwise_add_pattern](
+      const GraphPatternDetector::subgraph_t& subgraph)
+      -> std::tuple<Node*, Node*> {
+        GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_op, elementwise_add_op,
+                                  elementwise_add_pattern);
+        GET_IR_NODE_FROM_SUBGRAPH(elementwise_add_out, elementwise_add_out,
+                                  elementwise_add_pattern);
+
+        return std::make_tuple(elementwise_add_op, elementwise_add_out);
+      };
+
+  return ExecuteHandleOnGraph<ProjectionFuseHandle>(
+      &gpd, graph_with_stats,
+      [this,
+       &conv_x_pattern](const GraphPatternDetector::subgraph_t& subgraph) {
+        return GetNodesFromConv(conv_x_pattern, subgraph);
+      },
+      [this,
+       &conv_y_pattern](const GraphPatternDetector::subgraph_t& subgraph) {
+        return GetNodesFromConv(conv_y_pattern, subgraph);
+      },
+      get_node_from_elementwise_add);
+}
+
+graph_ptr ResidualConnectionMKLDNNFusePass::ApplyImpl(graph_ptr graph) const {
+  FusePassBase::Init(name_scope_, graph.get());
+  auto fused_graph_with_stats = FuseConvAsY(
+      name_scope_,
+      FuseConvAsX(
+          name_scope_,
+          FuseProjectionConv(name_scope_, std::make_pair(graph.get(), 0))));
+
+  std::cout << "Fused graph " << fused_graph_with_stats.second << std::endl;
+  AddStatis(fused_graph_with_stats.second);
   return graph;
 }
 }  // namespace ir
@@ -151,4 +435,4 @@ graph_ptr ConvElementwiseAddMKLDNNFusePass::ApplyImpl(graph_ptr graph) const {
 }  // namespace paddle
 
 REGISTER_PASS(conv_elementwise_add_mkldnn_fuse_pass,
-              paddle::framework::ir::ConvElementwiseAddMKLDNNFusePass);
+              paddle::framework::ir::ResidualConnectionMKLDNNFusePass);

paddle/fluid/framework/ir/conv_elementwise_add_mkldnn_fuse_pass.h

@@ -15,24 +15,119 @@
 #pragma once
 
 #include <string>
+#include <tuple>
+#include <utility>
 #include "paddle/fluid/framework/ir/fuse_pass_base.h"
 #include "paddle/fluid/framework/ir/graph.h"
 #include "paddle/fluid/framework/ir/graph_pattern_detector.h"
 
+#include <boost/optional.hpp>
+
 namespace paddle {
 namespace framework {
 namespace ir {
 
-class ConvElementwiseAddMKLDNNFusePass : public FusePassBase {
+using graph_ptr = std::unique_ptr<ir::Graph>;
+using GraphWithStats = std::pair<ir::Graph*, int>;
+
+void CorrectGraphEdges(Graph* graph, Node* from, Node* to);
+bool IsReachable(ir::Graph* graph, Node* from, Node* to);
+boost::optional<Node*> HasBias(const Node& op, const std::string& bias_name);
+
+class ResidualConnectionMKLDNNFusePass : public FusePassBase {
+ private:
+  GraphWithStats FuseConvAsX(const std::string& name_scope,
+                             const GraphWithStats& graph_with_stats) const;
+  GraphWithStats FuseConvAsY(const std::string& name_scope,
+                             const GraphWithStats& graph_with_stats) const;
+  GraphWithStats FuseProjectionConv(
+      const std::string& name_scope,
+      const GraphWithStats& graph_with_stats) const;
+
+  template <typename RetType>
+  using GetNodeFunc =
+      std::function<RetType(const GraphPatternDetector::subgraph_t& subgraph)>;
+  using IdentityConvFunc = GetNodeFunc<std::tuple<Node*, Node*, Node*, Node*>>;
+  using IdentityElementwiseAddFunc =
+      GetNodeFunc<std::tuple<Node*, Node*, Node*>>;
+
+  using ProjectionConvFunc = IdentityConvFunc;
+  using ProjectionElementwiseAddFunc = GetNodeFunc<std::tuple<Node*, Node*>>;
+
+  using CanFuseFunc = std::function<bool(Node*, Node*)>;
+
+  std::tuple<Node*, Node*, Node*, Node*> GetNodesFromConv(
+      const patterns::Conv& conv_pattern,
+      const GraphPatternDetector::subgraph_t& subgraph) const;
+
+  std::tuple<Node*, Node*, Node*, Node*> GetNodesFromProjectionConv(
+      const patterns::Conv& conv_pattern,
+      const GraphPatternDetector::subgraph_t& subgraph) const;
+
+  template <typename HandleType, typename... OpFuncs>
+  GraphWithStats ExecuteHandleOnGraph(GraphPatternDetector* gpd,
+                                      const GraphWithStats& graph_with_stats,
+                                      OpFuncs&&... op_funcs) const {
+    ir::Graph* graph;
+    int stats;
+
+    std::tie(graph, stats) = graph_with_stats;
+
+    auto can_fuse = [this](Node* op1, Node* op2) -> bool {
+      return this->FindFuseOption(*op1, *op2) == FUSE_MKLDNN;
+    };
+
+    auto fuse_handle = HandleType{can_fuse, std::forward<OpFuncs>(op_funcs)...};
+
+    (*gpd)(graph, fuse_handle);
+
+    return std::make_pair(graph, stats + fuse_handle.get_stats());
+  }
+
+  struct IdentityFuseHandle {
+    IdentityFuseHandle(
+        const CanFuseFunc& can_fuse_func,
+        const IdentityConvFunc& get_node_from_conv_op,
+        const IdentityElementwiseAddFunc& get_node_from_elementwise_add_op);
+
+    void operator()(const GraphPatternDetector::subgraph_t& subgraph,
+                    Graph* graph);
+    int get_stats() const { return *fusion_stats; }
+
+   private:
+    std::shared_ptr<int> fusion_stats;
+    CanFuseFunc can_fuse_func;
+    IdentityConvFunc get_node_from_conv_op;
+    IdentityElementwiseAddFunc get_node_from_elementwise_add_op;
+  };
+
+  struct ProjectionFuseHandle {
+    ProjectionFuseHandle(
+        const CanFuseFunc& can_fuse_func,
+        const ProjectionConvFunc& get_node_from_conv_x_op,
+        const ProjectionConvFunc& get_node_from_conv_y_op,
+        const ProjectionElementwiseAddFunc& get_node_from_elementwise_add_op);
+
+    void operator()(const GraphPatternDetector::subgraph_t& subgraph,
+                    Graph* graph);
+    int get_stats() const { return *fusion_stats; }
+
+   private:
+    std::shared_ptr<int> fusion_stats;
+    CanFuseFunc can_fuse_func;
+    ProjectionConvFunc get_node_from_conv_x_op;
+    ProjectionConvFunc get_node_from_conv_y_op;
+    ProjectionElementwiseAddFunc get_node_from_elementwise_add_op;
+  };
+
  public:
-  virtual ~ConvElementwiseAddMKLDNNFusePass() {}
+  virtual ~ResidualConnectionMKLDNNFusePass() {}
 
  protected:
-  std::unique_ptr<ir::Graph> ApplyImpl(std::unique_ptr<ir::Graph> graph) const;
+  std::unique_ptr<ir::Graph> ApplyImpl(graph_ptr graph) const;
 
-  const std::string name_scope_{"residual_connections_fuse_pass"};
+  const std::string name_scope_{"residual_connection_fuse_pass"};
 };
-
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/ir/conv_elementwise_add_mkldnn_fuse_pass_tester.cc

@@ -40,7 +40,7 @@ void SetOp(ProgramDesc* prog, const std::string& type,
   op->SetOutput(output.first, {output.second});
 }
 
-struct IsReachable {
+struct TestIsReachable {
   using func = std::function<bool(const std::string&, const std::string&)>;
 
   auto operator()(const std::unique_ptr<ir::Graph>& graph) -> func {
@@ -89,7 +89,9 @@ struct IsReachable {
   }
 };
 
-void AssertOpsCount(const std::unique_ptr<ir::Graph>& graph) {
+void AssertOpsCount(const std::unique_ptr<ir::Graph>& graph,
+                    int expected_conv_count,
+                    int expected_elementwise_add_count = 0) {
   int conv_count = 0;
   int elementwise_add_count = 0;
 
@@ -101,8 +103,8 @@ void AssertOpsCount(const std::unique_ptr<ir::Graph>& graph) {
       ++elementwise_add_count;
     }
   }
-  EXPECT_EQ(conv_count, 1);
-  EXPECT_EQ(elementwise_add_count, 0);
+  EXPECT_EQ(conv_count, expected_conv_count);
+  EXPECT_EQ(elementwise_add_count, expected_elementwise_add_count);
 }
 
 ProgramDesc BuildProgramDesc(const std::vector<std::string>& transient_vars,
@@ -127,22 +129,13 @@ ProgramDesc BuildProgramDesc(const std::vector<std::string>& transient_vars,
 
   return prog;
 }
-}  // namespace
-
-TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionWithElementwiseAddRelu) {
-  auto prog =
-      BuildProgramDesc({"a", "b", "c", "d", "e", "f"}, {"bias", "weights"});
-
-  SetOp(&prog, "conv2d",
-        {{"Input", "a"}, {"Bias", "bias"}, {"Filter", "weights"}},
-        {"Output", "b"});
-  SetOp(&prog, "elementwise_add", {{"X", "b"}, {"Y", "c"}}, {"Out", "d"});
-  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+void RunPassAndAssert(ProgramDesc* prog, const std::string& from,
+                      const std::string& to, int expected_conv_num) {
+  std::unique_ptr<ir::Graph> graph(new ir::Graph(*prog));
 
-  IsReachable is_reachable;
-  EXPECT_TRUE(is_reachable(graph)("a", "relu"));
+  TestIsReachable is_reachable;
+  EXPECT_TRUE(is_reachable(graph)(from, to));
 
   auto pass =
       PassRegistry::Instance().Get("conv_elementwise_add_mkldnn_fuse_pass");
@@ -150,82 +143,87 @@ TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionWithElementwiseAddRelu) {
   graph = pass->Apply(std::move(graph));
   int current_nodes_num = graph->Nodes().size();
 
-  EXPECT_TRUE(is_reachable(graph)("a", "relu"));
+  EXPECT_TRUE(is_reachable(graph)(from, to));
 
   EXPECT_EQ(original_nodes_num - nodes_removed + nodes_added,
             current_nodes_num);
 
-  AssertOpsCount(graph);
+  AssertOpsCount(graph, expected_conv_num);
 }
+}  // namespace
 
-TEST(ConvElementwiseAddMKLDNNFusePass,
-     ConvolutionWithElementwiseAddReluNoBias) {
-  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"weights"});
-  SetOp(&prog, "conv2d", {{"Input", "a"}, {"Filter", "weights"}},
-        {"Output", "b"});
-  SetOp(&prog, "elementwise_add", {{"X", "b"}, {"Y", "c"}}, {"Out", "d"});
-  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
-
-  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionAsYWithElementwiseAddRelu) {
+  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"bias", "weights"});
 
-  IsReachable is_reachable;
+  SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
+  SetOp(&prog, "conv2d",
+        {{"Input", "b"}, {"Bias", "bias"}, {"Filter", "weights"}},
+        {"Output", "c"});
 
-  EXPECT_TRUE(is_reachable(graph)("a", "relu"));
+  SetOp(&prog, "elementwise_add", {{"X", "a"}, {"Y", "c"}}, {"Out", "d"});
+  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  auto pass =
-      PassRegistry::Instance().Get("conv_elementwise_add_mkldnn_fuse_pass");
-  int original_nodes_num = graph->Nodes().size();
-  graph = pass->Apply(std::move(graph));
-  int current_nodes_num = graph->Nodes().size();
+  RunPassAndAssert(&prog, "a", "relu", 1);
+}
 
-  EXPECT_TRUE(is_reachable(graph)("a", "relu"));
+TEST(ConvElementwiseAddMKLDNNFusePass,
+     ConvolutionAsYWithElementwiseAddReluNoBias) {
+  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"weights"});
 
-  EXPECT_EQ(original_nodes_num - nodes_removed + nodes_added,
-            current_nodes_num);
+  SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
+  SetOp(&prog, "conv2d", {{"Input", "b"}, {"Filter", "weights"}},
+        {"Output", "c"});
+  SetOp(&prog, "elementwise_add", {{"X", "a"}, {"Y", "c"}}, {"Out", "d"});
+  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  AssertOpsCount(graph);
+  RunPassAndAssert(&prog, "a", "relu", 1);
 }
 
-TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionElementwiseAdd) {
-  auto prog = BuildProgramDesc({"a", "b", "c", "d"}, {"bias", "weights"});
+TEST(ConvElementwiseAddMKLDNNFusePass, ConvolutionAsXWithElementwiseAddRelu) {
+  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"bias", "weights"});
+
+  SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
   SetOp(&prog, "conv2d",
-        {{"Input", "a"}, {"Bias", "bias"}, {"Filter", "weights"}},
-        {"Output", "b"});
-  SetOp(&prog, "elementwise_add", {{"X", "b"}, {"Y", "c"}}, {"Out", "d"});
+        {{"Input", "b"}, {"Bias", "bias"}, {"Filter", "weights"}},
+        {"Output", "c"});
 
-  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+  SetOp(&prog, "elementwise_add", {{"X", "c"}, {"Y", "a"}}, {"Out", "d"});
+  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  IsReachable is_reachable;
-  EXPECT_TRUE(is_reachable(graph)("a", "d"));
+  RunPassAndAssert(&prog, "a", "relu", 1);
+}
 
-  auto pass =
-      PassRegistry::Instance().Get("conv_elementwise_add_mkldnn_fuse_pass");
-  int original_nodes_num = graph->Nodes().size();
-  graph = pass->Apply(std::move(graph));
-  int current_nodes_num = graph->Nodes().size();
+TEST(ConvElementwiseAddMKLDNNFusePass,
+     ConvolutionAsXWithElementwiseAddReluNoBias) {
+  auto prog = BuildProgramDesc({"a", "b", "c", "d", "e"}, {"weights"});
 
-  EXPECT_FALSE(is_reachable(graph)("a", "d"));
+  SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
+  SetOp(&prog, "conv2d", {{"Input", "b"}, {"Filter", "weights"}},
+        {"Output", "c"});
+  SetOp(&prog, "elementwise_add", {{"X", "c"}, {"Y", "a"}}, {"Out", "d"});
+  SetOp(&prog, "relu", {{"X", "d"}}, {"Out", "e"});
 
-  EXPECT_EQ(original_nodes_num - nodes_removed + nodes_added,
-            current_nodes_num);
-  AssertOpsCount(graph);
+  RunPassAndAssert(&prog, "a", "relu", 1);
 }
 
-TEST(ConvElementwiseAddMKLDNNFusePass, SigmoidConvolutionAddElementwiseRelu) {
+TEST(ConvElementwiseAddMKLDNNFusePass, NoFusion) {
   auto prog =
-      BuildProgramDesc({"a", "b", "c", "d", "e", "f"}, {"bias", "weights"});
+      BuildProgramDesc({"a", "b", "c", "d", "e", "f", "g"}, {"weights"});
+
   SetOp(&prog, "sigmoid", {{"X", "a"}}, {"Out", "b"});
-  SetOp(&prog, "conv2d",
-        {{"Input", "b"}, {"Bias", "bias"}, {"Filter", "weights"}},
+  SetOp(&prog, "conv2d", {{"Input", "b"}, {"Filter", "weights"}},
         {"Output", "c"});
-  SetOp(&prog, "elementwise_add", {{"X", "c"}, {"Y", "d"}}, {"Out", "e"});
-  SetOp(&prog, "relu", {{"X", "e"}}, {"Out", "f"});
 
-  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+  SetOp(&prog, "conv2d", {{"Input", "d"}, {"Filter", "weights"}},
+        {"Output", "e"});
 
-  IsReachable is_reachable;
+  SetOp(&prog, "elementwise_add", {{"X", "c"}, {"Y", "e"}}, {"Out", "f"});
+  SetOp(&prog, "relu", {{"X", "f"}}, {"Out", "g"});
 
-  EXPECT_TRUE(is_reachable(graph)("a", "f"));
+  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+
+  TestIsReachable is_reachable;
+  EXPECT_TRUE(is_reachable(graph)("a", "g"));
 
   auto pass =
       PassRegistry::Instance().Get("conv_elementwise_add_mkldnn_fuse_pass");
@@ -233,11 +231,10 @@ TEST(ConvElementwiseAddMKLDNNFusePass, SigmoidConvolutionAddElementwiseRelu) {
   graph = pass->Apply(std::move(graph));
   int current_nodes_num = graph->Nodes().size();
 
-  EXPECT_TRUE(is_reachable(graph)("a", "f"));
+  EXPECT_TRUE(is_reachable(graph)("a", "g"));
+  EXPECT_EQ(original_nodes_num, current_nodes_num);
 
-  EXPECT_EQ(original_nodes_num - nodes_removed + nodes_added,
-            current_nodes_num);
-  AssertOpsCount(graph);
+  AssertOpsCount(graph, 2, 1);
 }
 
 }  // namespace ir

paddle/fluid/framework/ir/graph_helper.cc

@@ -15,8 +15,15 @@ limitations under the License. */
 #include "paddle/fluid/framework/ir/graph_helper.h"
 #include <algorithm>
 #include <deque>
+#include <fstream>
+#include <iosfwd>
+#include <ostream>
 #include <unordered_set>
 
+DEFINE_string(print_sub_graph_dir, "",
+              "FLAGS_print_sub_graph_dir is used "
+              "to print the nodes of sub_graphs.");
+
 namespace paddle {
 namespace framework {
 namespace ir {
@@ -164,12 +171,15 @@ size_t GraphNum(const Graph &graph) {
     graph_nodes.emplace_back(g_nodes);
   }
 
-  if (VLOG_IS_ON(100)) {
-    VLOG(100) << "graph_num: " << graph_nodes.size();
-    for (auto &g_n : graph_nodes) {
-      VLOG(100) << "graph_nodes: " << g_n.size();
-      if (g_n.size() < 10) {
-        std::stringstream out;
+  if (FLAGS_print_sub_graph_dir.size()) {
+    if (graph_nodes.size() > 1) {
+      std::stringstream out;
+      for (auto &g_n : graph_nodes) {
+        out << "graph_nodes: " << g_n.size() << "\n";
+      }
+      out << "\n\n";
+      for (auto &g_n : graph_nodes) {
+        out << "graph_nodes: " << g_n.size();
         for (auto &node : g_n) {
           out << "\nNode: " << node->Name() << " in [";
           for (auto &n : node->inputs) {
@@ -181,8 +191,12 @@ size_t GraphNum(const Graph &graph) {
           }
           out << "]";
         }
-        VLOG(100) << out.str();
+        out << "\n\n\n";
       }
+      std::unique_ptr<std::ostream> fout(
+          new std::ofstream(FLAGS_print_sub_graph_dir));
+      PADDLE_ENFORCE(fout->good());
+      *fout << out.str();
     }
   }
 

paddle/fluid/framework/ir/graph_pattern_detector.cc

@@ -1084,16 +1084,12 @@ PDNode *patterns::Conv::operator()() {
   return output_var;
 }
 
-PDNode *patterns::ElementwiseAdd::operator()(PDNode *x_var) {
+PDNode *patterns::ElementwiseAdd::operator()(PDNode *x_var, PDNode *y_var) {
   auto elementwise_add_op = pattern->NewNode(elementwise_add_op_repr())
                                 ->assert_is_op("elementwise_add");
 
-  x_var->assert_is_op_input("elementwise_add", "X");
-
-  auto y_var = pattern->NewNode(elementwise_add_x_repr())
-                   ->AsInput()
-                   ->assert_is_op_input("elementwise_add", "Y");
-
+  x_var->AsInput()->assert_is_op_input("elementwise_add", "X");
+  y_var->AsInput()->assert_is_op_input("elementwise_add", "Y");
   auto out_var = pattern->NewNode(elementwise_add_out_repr())
                      ->AsOutput()
                      ->assert_is_op_output("elementwise_add", "Out");

paddle/fluid/framework/ir/graph_pattern_detector.h

@@ -664,7 +664,7 @@ struct ElementwiseAdd : public PatternBase {
   ElementwiseAdd(PDPattern* pattern, const std::string& name_scope)
       : PatternBase(pattern, name_scope, "elementwise_add") {}
 
-  PDNode* operator()(PDNode* x_var);
+  PDNode* operator()(PDNode* x_var, PDNode* y_var);
 
   PATTERN_DECL_NODE(elementwise_add_op);
   PATTERN_DECL_NODE(elementwise_add_x);

paddle/fluid/framework/lod_tensor.h

@@ -111,9 +111,6 @@ class LoDTensor : public Tensor {
  public:
   LoDTensor() : Tensor() {}
 
-  /* Constructor with place should only be used in pybind */
-  explicit LoDTensor(const platform::Place& place) : Tensor(place) {}
-
   explicit LoDTensor(const LoD& lod) : lod_(lod) {}
 
   void set_lod(const LoD& lod) { lod_ = lod; }

paddle/fluid/framework/mixed_vector.h

@@ -23,6 +23,7 @@
 #include "paddle/fluid/framework/details/cow_ptr.h"
 #include "paddle/fluid/framework/tensor.h"
 #include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/memory/malloc.h"
 #include "paddle/fluid/memory/memcpy.h"
 
 #include "glog/logging.h"
@@ -31,46 +32,6 @@ namespace paddle {
 namespace framework {
 
 #if defined(PADDLE_WITH_CUDA)
-namespace details {
-struct CUDABuffer {
-  void *data_{nullptr};
-  size_t size_{0};
-  platform::CUDAPlace place_;
-
-  CUDABuffer() {}
-  CUDABuffer(platform::Place place, size_t size)
-      : size_(size), place_(boost::get<platform::CUDAPlace>(place)) {
-    data_ = memory::Alloc(place_, size);
-  }
-
-  ~CUDABuffer() { ClearMemory(); }
-
-  CUDABuffer(const CUDABuffer &o) = delete;
-  CUDABuffer &operator=(const CUDABuffer &o) = delete;
-
-  void Resize(platform::Place place, size_t size) {
-    ClearMemory();
-    place_ = boost::get<platform::CUDAPlace>(place);
-    data_ = memory::Alloc(place_, size);
-    PADDLE_ENFORCE_NOT_NULL(data_);
-    size_ = size;
-  }
-
-  void Swap(CUDABuffer &o) {
-    std::swap(data_, o.data_);
-    std::swap(place_, o.place_);
-    std::swap(size_, o.size_);
-  }
-
- private:
-  void ClearMemory() const {
-    if (data_ != nullptr) {
-      memory::Free(place_, data_);
-    }
-  }
-};
-}  // namespace details
-
 // Vector<T> implements the std::vector interface, and can get Data or
 // MutableData from any place. The data will be synced implicitly inside.
 template <typename T>
@@ -103,8 +64,6 @@ class Vector {
       o.ImmutableCPU();
       cpu_ = o.cpu_;
       flag_ = kDataInCPU;
-      details::CUDABuffer null;
-      gpu_.Swap(null);
       return *this;
     }
 
@@ -199,7 +158,7 @@ class Vector {
       PADDLE_ENFORCE(platform::is_gpu_place(place),
                      "CUDA Data must on CUDA place");
       ImmutableCUDA(place);
-      return reinterpret_cast<T *>(gpu_.data_);
+      return reinterpret_cast<T *>(gpu_->ptr());
     }
 
     // get cuda ptr. mutable
@@ -234,13 +193,11 @@ class Vector {
 
     std::mutex &Mutex() const { return mtx_; }
 
-    std::unique_ptr<platform::CUDAPlace> CUDAPlace() const {
-      if (gpu_.data_ == nullptr) {
-        return nullptr;
-      } else {
-        return std::unique_ptr<platform::CUDAPlace>(
-            new platform::CUDAPlace(gpu_.place_));
-      }
+    boost::optional<platform::CUDAPlace> CUDAPlace() const {
+      return gpu_ == nullptr
+                 ? boost::none
+                 : boost::optional<platform::CUDAPlace>(
+                       boost::get<platform::CUDAPlace>(gpu_->place()));
     }
 
    private:
@@ -254,13 +211,12 @@ class Vector {
     void CopyToCPU() const {
       // COPY GPU Data To CPU
       auto *dev_ctx = static_cast<platform::CUDADeviceContext *>(
-          platform::DeviceContextPool::Instance().Get(
-              platform::Place(gpu_.place_)));
+          platform::DeviceContextPool::Instance().Get(gpu_->place()));
       auto stream = dev_ctx->stream();
-      void *src = gpu_.data_;
+      void *src = gpu_->ptr();
       void *dst = cpu_.data();
-      memory::Copy(platform::CPUPlace(), dst, gpu_.place_, src, gpu_.size_,
-                   stream);
+      memory::Copy(platform::CPUPlace(), dst, CUDAPlace().get(), src,
+                   gpu_->size(), stream);
       dev_ctx->Wait();
     }
 
@@ -277,8 +233,7 @@ class Vector {
           CopyCPUDataToCUDA(place);
           UnsetFlag(kDirty);
           SetFlag(kDataInCUDA);
-        } else if (IsInCUDA() &&
-                   !(boost::get<platform::CUDAPlace>(place) == gpu_.place_)) {
+        } else if (IsInCUDA() && !(place == gpu_->place())) {
           PADDLE_THROW("This situation should not happen");
           // Still dirty
         } else {
@@ -290,7 +245,7 @@ class Vector {
           // Even data is not dirty. However, data is not in CUDA. Copy data.
           CopyCPUDataToCUDA(place);
           SetFlag(kDataInCUDA);
-        } else if (!(boost::get<platform::CUDAPlace>(place) == gpu_.place_)) {
+        } else if (!(place == gpu_->place())) {
           PADDLE_THROW("This situation should not happen.");
         } else {
           // Not Dirty && DataInCUDA && Device is same
@@ -301,13 +256,13 @@ class Vector {
 
     void CopyCPUDataToCUDA(const platform::Place &place) const {
       void *src = cpu_.data();
-      gpu_.Resize(place, cpu_.size() * sizeof(T));
-      void *dst = gpu_.data_;
+      gpu_ = memory::Alloc(place, cpu_.size() * sizeof(T));
+      void *dst = gpu_->ptr();
       auto *dev_ctx = static_cast<platform::CUDADeviceContext *>(
           platform::DeviceContextPool::Instance().Get(place));
       auto stream = dev_ctx->stream();
-      memory::Copy(gpu_.place_, dst, platform::CPUPlace(), src, gpu_.size_,
-                   stream);
+      memory::Copy(CUDAPlace().get(), dst, platform::CPUPlace(), src,
+                   gpu_->size(), stream);
     }
 
     void ImmutableCPU() const {
@@ -329,7 +284,7 @@ class Vector {
     bool IsInCPU() const { return flag_ & kDataInCPU; }
 
     mutable std::vector<T> cpu_;
-    mutable details::CUDABuffer gpu_;
+    mutable memory::AllocationPtr gpu_;
     mutable int flag_;
 
     mutable std::mutex mtx_;
@@ -428,8 +383,8 @@ class Vector {
       auto &mtx = m_.Data().Mutex();
       std::lock_guard<std::mutex> guard(mtx);
       auto cuda_place = m_.Data().CUDAPlace();
-      if (cuda_place == nullptr ||
-          *cuda_place == boost::get<platform::CUDAPlace>(place)) {
+      if (cuda_place == boost::none ||
+          cuda_place == boost::get<platform::CUDAPlace>(place)) {
         return m_.Data().CUDAData(place);
       }
     }
@@ -444,8 +399,8 @@ class Vector {
       auto &mtx = m_.Data().Mutex();
       std::lock_guard<std::mutex> guard(mtx);
       auto cuda_place = m_.Data().CUDAPlace();
-      if (cuda_place == nullptr ||
-          *cuda_place == boost::get<platform::CUDAPlace>(place)) {
+      if (cuda_place == boost::none ||
+          cuda_place == boost::get<platform::CUDAPlace>(place)) {
         return m_.MutableData()->CUDAMutableData(place);
       }
     }

paddle/fluid/framework/parallel_executor.cc

@@ -171,8 +171,17 @@ ParallelExecutor::ParallelExecutor(
   }
   // If the loss_var_name is given, the number of graph should be only one.
   if (loss_var_name.size()) {
-    PADDLE_ENFORCE_EQ(ir::GraphNum(*graph), 1,
-                      "The number of graph should be only one");
+    size_t graph_num = ir::GraphNum(*graph);
+    if (graph_num > 1) {
+      LOG(WARNING)
+          << "The number of graph should be only one, "
+             "but the current graph has "
+          << ir::GraphNum(*graph)
+          << " sub_graphs. If you want to see the nodes of the "
+             "sub_graphs, you should use 'FLAGS_print_sub_graph_dir' "
+             "to specify the output dir. NOTES: if you not do training, "
+             "please don't pass loss_var_name.";
+    }
   }
 
   if (exec_strategy.type_ == ExecutionStrategy::kDefault) {

paddle/fluid/framework/tensor.cc

@@ -32,10 +32,9 @@ size_t Tensor::memory_size() const {
 }
 
 void* Tensor::mutable_data(platform::Place place, std::type_index type,
+                           memory::Allocator::Attr attr,
                            size_t requested_size) {
-  if (holder_ != nullptr) {
-    holder_->set_type(type);
-  }
+  type_ = type;
   PADDLE_ENFORCE_GE(numel(), 0,
                     "When calling this method, the Tensor's numel must be "
                     "equal or larger than zero. "
@@ -48,35 +47,18 @@ void* Tensor::mutable_data(platform::Place place, std::type_index type,
   /* some versions of boost::variant don't have operator!= */
   if (holder_ == nullptr || !(holder_->place() == place) ||
       holder_->size() < size + offset_) {
-    if (platform::is_cpu_place(place)) {
-      holder_.reset(new PlaceholderImpl<platform::CPUPlace>(
-          boost::get<platform::CPUPlace>(place), size, type));
-    } else if (platform::is_gpu_place(place) ||
-               platform::is_cuda_pinned_place(place)) {
-#ifndef PADDLE_WITH_CUDA
-      PADDLE_THROW(
-          "CUDAPlace or CUDAPinnedPlace is not supported in CPU-only mode.");
-    }
-#else
-      if (platform::is_gpu_place(place)) {
-        holder_.reset(new PlaceholderImpl<platform::CUDAPlace>(
-            boost::get<platform::CUDAPlace>(place), size, type));
-      } else if (platform::is_cuda_pinned_place(place)) {
-        holder_.reset(new PlaceholderImpl<platform::CUDAPinnedPlace>(
-            boost::get<platform::CUDAPinnedPlace>(place), size, type));
-      }
-    }
-#endif
+    holder_ = memory::AllocShared(place, size, attr);
     offset_ = 0;
   }
   return reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(holder_->ptr()) +
                                  offset_);
 }
 
-void* Tensor::mutable_data(platform::Place place, size_t requested_size) {
+void* Tensor::mutable_data(platform::Place place, memory::Allocator::Attr attr,
+                           size_t requested_size) {
   PADDLE_ENFORCE(this->holder_ != nullptr,
                  "Cannot invoke mutable data if current hold nothing.");
-  return mutable_data(place, holder_->type(), requested_size);
+  return mutable_data(place, type_, attr, requested_size);
 }
 
 Tensor& Tensor::ShareDataWith(const Tensor& src) {
@@ -101,6 +83,7 @@ Tensor Tensor::Slice(int begin_idx, int end_idx) const {
     Tensor dst;
     dst.holder_ = holder_;
     dst.set_layout(layout_);
+    dst.type_ = type_;
     DDim dst_dims = dims_;
     dst_dims[0] = end_idx - begin_idx;
     dst.Resize(dst_dims);

paddle/fluid/framework/tensor.h

@@ -67,12 +67,7 @@ class Tensor {
   friend struct EigenVector;
 
  public:
-  Tensor() : offset_(0) {}
-
-  /*! Constructor with place should only be used in pybind. */
-  explicit Tensor(const platform::Place& place) : offset_(0) {
-    holder_->set_place(place);
-  }
+  Tensor() : type_(typeid(float)), offset_(0) {}
 
   /*! Return a pointer to mutable memory block. */
   template <typename T>
@@ -89,12 +84,17 @@ class Tensor {
    * @note    If not exist, then allocation.
    */
   template <typename T>
-  T* mutable_data(platform::Place place, size_t requested_size = 0);
+  T* mutable_data(platform::Place place,
+                  memory::Allocator::Attr attr = memory::Allocator::kDefault,
+                  size_t requested_size = 0);
 
   void* mutable_data(platform::Place place, std::type_index type,
+                     memory::Allocator::Attr attr = memory::Allocator::kDefault,
                      size_t requested_size = 0);
 
-  void* mutable_data(platform::Place place, size_t requested_size = 0);
+  void* mutable_data(platform::Place place,
+                     memory::Allocator::Attr attr = memory::Allocator::kDefault,
+                     size_t requested_size = 0);
 
   /**
    * @brief     Return a pointer to mutable memory block.
@@ -106,7 +106,9 @@ class Tensor {
    * @note      If not exist, then allocation.
    */
   template <typename T>
-  T* mutable_data(DDim dims, platform::Place place, size_t requested_size = 0);
+  T* mutable_data(DDim dims, platform::Place place,
+                  memory::Allocator::Attr attr = memory::Allocator::kDefault,
+                  size_t requested_size = 0);
 
   /*! Return the dimensions of the memory block. */
   const DDim& dims() const;
@@ -139,7 +141,7 @@ class Tensor {
   std::type_index type() const {
     PADDLE_ENFORCE_NOT_NULL(
         holder_, "Tensor not initialized yet when Tensor::type() is called.");
-    return holder_->type();
+    return type_;
   }
 
   // memory size returns the holding memory size in byte.
@@ -153,56 +155,13 @@ class Tensor {
 
   void clear() { holder_ = nullptr; }
 
- private:
-  /**
-   * @note    Placeholder hides type T, so it doesn't appear as a template
-   *          parameter of Variable.
-   */
-  struct Placeholder {
-    virtual ~Placeholder() = default;
-    virtual void* ptr() const = 0;
-    virtual size_t size() const = 0;
-    virtual std::type_index type() const = 0;
-    virtual platform::Place place() const = 0;
-    virtual void set_type(std::type_index type) = 0;
-    virtual void set_place(platform::Place place) = 0;
-  };
-
-  template <typename Place>
-  struct PlaceholderImpl : public Placeholder {
-    PlaceholderImpl(Place place, size_t size, std::type_index type)
-        : ptr_(static_cast<uint8_t*>(memory::Alloc(place, size)),
-               memory::PODDeleter<uint8_t, Place>(place)),
-          place_(place),
-          size_(size),
-          type_(type) {
-      PADDLE_ENFORCE_NOT_NULL(ptr_, "Insufficient %s memory to allocation.",
-                              (is_cpu_place(place_) ? "CPU" : "GPU"));
-    }
-
-    virtual size_t size() const { return size_; }
-    virtual platform::Place place() const { return place_; }
-    virtual void* ptr() const { return static_cast<void*>(ptr_.get()); }
-    virtual std::type_index type() const { return type_; }
-    virtual void set_type(std::type_index type) { type_ = type; }
-    virtual void set_place(platform::Place place) { place_ = place; }
-
-    /*! the pointer of memory block. */
-    std::unique_ptr<uint8_t, memory::PODDeleter<uint8_t, Place>> ptr_;
-
-    /*! the place of memory block. */
-    platform::Place place_;
-
-    /*! the size of memory block. */
-    size_t size_;
-
-    /* the current type of memory */
-    std::type_index type_;
-  };
+  const std::shared_ptr<memory::Allocation>& Holder() const { return holder_; }
+  size_t offset() const { return offset_; }
 
+ private:
   /*! holds the memory block if allocated. */
-  std::shared_ptr<Placeholder> holder_;
-
+  std::shared_ptr<memory::Allocation> holder_;
+  std::type_index type_;
   /**
    * @brief points to elements dimensions.
    *

paddle/fluid/framework/tensor_impl.h

@@ -23,10 +23,10 @@ namespace framework {
 template <typename T>
 inline const T* Tensor::data() const {
   check_memory_size();
-  bool valid = std::is_same<T, void>::value ||
-               holder_->type() == std::type_index(typeid(T));
+  bool valid =
+      std::is_same<T, void>::value || type_ == std::type_index(typeid(T));
   PADDLE_ENFORCE(valid, "Tensor holds the wrong type, it holds %s",
-                 this->holder_->type().name());
+                 type_.name());
 
   return reinterpret_cast<const T*>(
       reinterpret_cast<uintptr_t>(holder_->ptr()) + offset_);
@@ -37,26 +37,30 @@ inline bool Tensor::IsInitialized() const { return holder_ != nullptr; }
 template <typename T>
 inline T* Tensor::data() {
   check_memory_size();
-  bool valid = std::is_same<T, void>::value ||
-               holder_->type() == std::type_index(typeid(T));
+  bool valid =
+      std::is_same<T, void>::value || type_ == std::type_index(typeid(T));
   PADDLE_ENFORCE(valid, "Tensor holds the wrong type, it holds %s",
-                 this->holder_->type().name());
+                 type_.name());
   return reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(holder_->ptr()) +
                               offset_);
 }
 
 template <typename T>
 inline T* Tensor::mutable_data(DDim dims, platform::Place place,
+                               memory::Allocator::Attr attr,
                                size_t requested_size) {
   static_assert(std::is_pod<T>::value, "T must be POD");
   Resize(dims);
-  return mutable_data<T>(place, requested_size);
+  return mutable_data<T>(place, attr, requested_size);
 }
 
 template <typename T>
-inline T* Tensor::mutable_data(platform::Place place, size_t requested_size) {
+inline T* Tensor::mutable_data(platform::Place place,
+                               memory::Allocator::Attr attr,
+                               size_t requested_size) {
   static_assert(std::is_pod<T>::value, "T must be POD");
-  return reinterpret_cast<T*>(mutable_data(place, typeid(T), requested_size));
+  return reinterpret_cast<T*>(
+      mutable_data(place, typeid(T), attr, requested_size));
 }
 
 inline Tensor ReshapeToMatrix(const Tensor& src, int num_col_dims) {

paddle/fluid/framework/tensor_util_test.cc

@@ -379,7 +379,9 @@ TEST(Tensor, FromAndToStream) {
     TensorToStream(oss, gpu_tensor, gpu_ctx);
 
     std::istringstream iss(oss.str());
-    TensorFromStream(iss, &dst_tensor, gpu_ctx);
+    TensorFromStream(
+        iss, &dst_tensor,
+        *platform::DeviceContextPool::Instance().Get(platform::CPUPlace()));
 
     int* dst_ptr = dst_tensor.mutable_data<int>(platform::CPUPlace());
     for (int i = 0; i < 6; ++i) {

paddle/fluid/inference/analysis/ir_passes/tensorrt_subgraph_pass.cc

@@ -114,7 +114,7 @@ void TensorRtSubgraphPass::CreateTensorRTOp(framework::ir::Node *node,
   // it is either an OP's input or an OP's output.
 
   auto &subgraph_nodes = *Agent(node).subgraph();
-  for (size_t index = 0; index < block_desc.OpSize(); index++) {
+  for (size_t index = 0; index < block_desc.OpSize(); ++index) {
     framework::proto::OpDesc *op = block_desc.Op(index)->Proto();
     auto correspond_node = subgraph_nodes[index];
     PADDLE_ENFORCE_EQ(correspond_node->Name(), op->type());

paddle/fluid/inference/analysis/passes/ir_analysis_compose_pass.cc

@@ -45,7 +45,8 @@ void IrAnalysisComposePass::InitTensorRTAttrs(Argument *argument) {
       std::unordered_set<std::string> teller_set(
           {"mul", "conv2d", "pool2d", "relu", "softmax", "sigmoid",
            "depthwise_conv2d", "batch_norm", "concat", "tanh", "pad",
-           "elementwise_add", "dropout", "split", "prelu", "conv2d_transpose"});
+           "elementwise_add", "elementwise_mul", "dropout", "split", "prelu",
+           "conv2d_transpose"});
       if (!node->IsOp()) return false;
 
       if (teller_set.count(node->Op()->Type())) {

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

 # Add TRT tests
 nv_library(tensorrt_converter
-  SRCS mul_op.cc conv2d_op.cc fc_op.cc pool2d_op.cc elementwise_op.cc
-batch_norm_op.cc activation_op.cc softmax_op.cc concat_op.cc dropout_op.cc
-pad_op.cc split_op.cc prelu_op.cc
-  DEPS tensorrt_engine tensorrt_plugin operator scope framework_proto op_registry)
+           SRCS mul_op.cc conv2d_op.cc fc_op.cc pool2d_op.cc elementwise_op.cc
+                batch_norm_op.cc activation_op.cc softmax_op.cc concat_op.cc dropout_op.cc
+                pad_op.cc split_op.cc prelu_op.cc
+           DEPS tensorrt_engine tensorrt_plugin operator scope framework_proto op_registry)
 
 nv_test(test_op_converter SRCS test_op_converter.cc DEPS
   ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine tensorrt_converter)
@@ -20,7 +20,8 @@ nv_test(test_trt_conv_op SRCS test_conv2d_op.cc conv2d_op.cc
 nv_test(test_trt_pool2d_op SRCS test_pool2d_op.cc pool2d_op.cc
         DEPS ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine pool_op SERIAL)
 nv_test(test_trt_elementwise_op SRCS test_elementwise_op.cc elementwise_op.cc
-        DEPS ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine elementwise_add_op SERIAL)
+        DEPS ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine tensorrt_plugin
+             elementwise_add_op elementwise_mul_op SERIAL)
 nv_test(test_trt_softmax_op SRCS test_softmax_op.cc softmax_op.cc
         DEPS ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine softmax_op SERIAL)
 nv_test(test_trt_batch_norm_op SRCS test_batch_norm_op.cc batch_norm_op.cc
@@ -33,7 +34,7 @@ nv_test(test_trt_pad_op SRCS test_pad_op.cc pad_op.cc
         DEPS ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine pad_op SERIAL)
 nv_test(test_trt_split_op SRCS test_split_op.cc split_op.cc
         DEPS ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine tensorrt_plugin
-        split_op concat_op SERIAL)
+             split_op concat_op SERIAL)
 nv_test(test_trt_prelu_op SRCS test_prelu_op.cc prelu_op.cc
         DEPS ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine tensorrt_plugin
         prelu_op SERIAL)

paddle/fluid/inference/tensorrt/convert/elementwise_op.cc

@@ -4,7 +4,7 @@ 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
+    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,
@@ -13,11 +13,25 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+#include "paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.h"
 
 namespace paddle {
 namespace inference {
 namespace tensorrt {
 
+static bool CheckDims(const nvinfer1::Dims& dims_x,
+                      const nvinfer1::Dims& dims_y) {
+  if (dims_x.nbDims != dims_y.nbDims) {
+    return false;
+  }
+  for (int i = 0; i < dims_x.nbDims; i++) {
+    if (dims_x.d[i] != dims_y.d[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
 class ElementwiseWeightOpConverter : public OpConverter {
  public:
   ElementwiseWeightOpConverter() {}
@@ -26,7 +40,7 @@ class ElementwiseWeightOpConverter : public OpConverter {
     // Here the two nullptr looks strange, that's because the
     // framework::OpDesc's constructor is strange.
     framework::OpDesc op_desc(op, nullptr);
-    VLOG(3) << "convert a fluid elementwise op to tensorrt IScaleLayer";
+    VLOG(3) << "Convert a fluid elementwise op to TensorRT IScaleLayer";
 
     PADDLE_ENFORCE_EQ(op_desc.Input("X").size(), 1);
     PADDLE_ENFORCE_EQ(op_desc.Input("Y").size(), 1);  // Y is a weight
@@ -106,10 +120,12 @@ class ElementwiseTensorOpConverter : public OpConverter {
   ElementwiseTensorOpConverter() {}
   void operator()(const framework::proto::OpDesc& op,
                   const framework::Scope& scope, bool test_mode) override {
+    auto op_pair = ops.find(op_type_);
+    PADDLE_ENFORCE(op_pair != ops.end(), "Wrong elementwise op type!");
+
     // Here the two nullptr looks strange, that's because the
     // framework::OpDesc's constructor is strange.
     framework::OpDesc op_desc(op, nullptr);
-    VLOG(3) << "convert a fluid elementwise op to tensorrt IScaleLayer";
 
     PADDLE_ENFORCE_EQ(op_desc.Input("X").size(), 1);
     PADDLE_ENFORCE_EQ(op_desc.Input("Y").size(), 1);  // Y is a weight
@@ -120,29 +136,35 @@ class ElementwiseTensorOpConverter : public OpConverter {
     nvinfer1::Dims dims_x = X->getDimensions();
     nvinfer1::Dims dims_y = Y->getDimensions();
 
-    // The two input tensor should have the same dims
-    PADDLE_ENFORCE(dims_x.nbDims >= 3);
-    if (dims_x.nbDims == dims_y.nbDims) {
-      for (int i = 0; i < dims_x.nbDims; i++) {
-        if (dims_x.d[i] != dims_y.d[i])
-          PADDLE_THROW("TensorRT unsupported tensor shape for Elementwise op!");
-      }
-    } else {
-      PADDLE_THROW("TensorRT unsupported tensor shape for Elementwise op!");
-    }
+    int axis = boost::get<int>(op_desc.GetAttr("axis"));
+    auto output_name = op_desc.Output("Out")[0];
+    if (CheckDims(dims_x, dims_y)) {
+      // The two input tensor should have the same dims
+      VLOG(3) << "Convert a fluid elementwise op to TensorRT IElementWiseLayer";
 
-    auto op_pair = ops.find(op_type_);
-    if (op_pair == ops.end()) {
-      PADDLE_THROW("Wrong elementwise op type!");
-    }
-    nvinfer1::IElementWiseLayer* layer = TRT_ENGINE_ADD_LAYER(
-        engine_, ElementWise, *const_cast<nvinfer1::ITensor*>(X),
-        *const_cast<nvinfer1::ITensor*>(Y), op_pair->second);
+      nvinfer1::IElementWiseLayer* layer = TRT_ENGINE_ADD_LAYER(
+          engine_, ElementWise, *const_cast<nvinfer1::ITensor*>(X),
+          *const_cast<nvinfer1::ITensor*>(Y), op_pair->second);
 
-    auto output_name = op_desc.Output("Out")[0];
-    layer->setName(("elementwise (Output: " + output_name + ")").c_str());
-    layer->getOutput(0)->setName(output_name.c_str());
-    engine_->SetITensor(output_name, layer->getOutput(0));
+      layer->setName(("elementwise (Output: " + output_name + ")").c_str());
+      layer->getOutput(0)->setName(output_name.c_str());
+      engine_->SetITensor(output_name, layer->getOutput(0));
+    } else {
+      VLOG(3) << "Convert a fluid elementwise op to TensorRT "
+                 "ElementWisePluginLayer";
+
+      plugin::ElementWisePlugin* plugin =
+          new plugin::ElementWisePlugin(op_pair->second, dims_x, dims_y, axis);
+      plugin->AddInput(X);
+      plugin->AddInput(Y);
+      nvinfer1::IPluginLayer* layer = engine_->AddPlugin(
+          const_cast<nvinfer1::ITensor* const*>(plugin->GetInputs().data()), 2,
+          reinterpret_cast<plugin::PluginTensorRT*>(plugin));
+
+      layer->setName(("elementwise (Output: " + output_name + ")").c_str());
+      layer->getOutput(0)->setName(output_name.c_str());
+      engine_->SetITensor(output_name, layer->getOutput(0));
+    }
     if (test_mode) {  // the test framework can not determine which is the
                       // output, so place the declaration inside.
       engine_->DeclareOutput(output_name);

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

@@ -61,7 +61,7 @@ class OpConverter {
       // TODO(xingzhaolong): all mul, sub, div
       // static std::unordered_set<std::string> add_weight_op_set {"add", "mul",
       // "sub", "div"};
-      static std::unordered_set<std::string> add_weight_op_set{"add"};
+      static std::unordered_set<std::string> add_weight_op_set{"add", "mul"};
       PADDLE_ENFORCE_EQ(op_desc.Input("Y").size(), 1UL);
       int op_type_len = op_desc.Type().size();
       std::string op_type = op_desc.Type().substr(op_type_len - 3, op_type_len);

paddle/fluid/inference/tensorrt/convert/prelu_op.cc

@@ -54,7 +54,7 @@ class PReluOpConverter : public OpConverter {
     TensorRTEngine::Weight alpha_rt(nvinfer1::DataType::kFLOAT,
                                     static_cast<void*>(alpha_data),
                                     alpha_tensor_device->numel());
-    PReluPlugin* plugin = new PReluPlugin(alpha_rt, mode);
+    plugin::PReluPlugin* plugin = new plugin::PReluPlugin(alpha_rt, mode);
     nvinfer1::IPluginLayer* layer =
         engine_->AddPlugin(&input, input_num, plugin);
     // keep alpha tensor to avoid release it's memory

paddle/fluid/inference/tensorrt/convert/split_op.cc

@@ -50,7 +50,7 @@ class SplitOpConverter : public OpConverter {
     PADDLE_ENFORCE(output_lengths.size() == output_num);
 
     //
-    SplitPlugin* plugin = new SplitPlugin(axis, output_lengths);
+    plugin::SplitPlugin* plugin = new plugin::SplitPlugin(axis, output_lengths);
     nvinfer1::IPluginLayer* layer =
         engine_->AddPlugin(&input, input_num, plugin);
 

paddle/fluid/inference/tensorrt/convert/test_elementwise_op.cc

@@ -20,13 +20,12 @@ namespace paddle {
 namespace inference {
 namespace tensorrt {
 
-TEST(elementwise_op, add_weight_test) {
+TEST(elementwise_op, add_weight) {
   std::unordered_set<std::string> parameters({"elementwise_add-Y"});
   framework::Scope scope;
   TRTConvertValidation validator(10, parameters, scope, 1 << 15);
   validator.DeclInputVar("elementwise_add-X", nvinfer1::DimsCHW(10, 3, 3));
   validator.DeclParamVar("elementwise_add-Y", nvinfer1::Dims3(10, 1, 1));
-  // validator.DeclParamVar("mul-Y", nvinfer1::Dims2(8, 2));
   validator.DeclOutputVar("elementwise_add-Out", nvinfer1::DimsCHW(10, 3, 3));
 
   // Prepare Op description
@@ -44,30 +43,65 @@ TEST(elementwise_op, add_weight_test) {
   validator.Execute(8);
 }
 
-TEST(elementwise_op, add_tensor_test) {
-  std::unordered_set<std::string> parameters;
-  framework::Scope scope;
-  TRTConvertValidation validator(8, parameters, scope, 1 << 15);
-  validator.DeclInputVar("elementwise_add-X", nvinfer1::DimsCHW(10, 3, 3));
-  validator.DeclInputVar("elementwise_add-Y", nvinfer1::Dims3(10, 3, 3));
-  // validator.DeclParamVar("mul-Y", nvinfer1::Dims2(8, 2));
-  validator.DeclOutputVar("elementwise_add-Out", nvinfer1::DimsCHW(10, 3, 3));
-
-  // Prepare Op description
-  framework::OpDesc desc;
-  desc.SetType("elementwise_add");
-  desc.SetInput("X", {"elementwise_add-X"});
-  desc.SetInput("Y", {"elementwise_add-Y"});
-  desc.SetOutput("Out", {"elementwise_add-Out"});
-
-  // the defalut axis of elementwise op is -1
-
-  validator.SetOp(*desc.Proto());
+TEST(elementwise_op, native) {
+  for (std::string type : {"add", "mul"}) {
+    int batch_size = 8;
+    std::unordered_set<std::string> parameters;
+    framework::Scope scope;
+    TRTConvertValidation validator(batch_size, parameters, scope, 1 << 15);
+    validator.DeclInputVar("elementwise_" + type + "-X",
+                           nvinfer1::DimsCHW(10, 3, 3));
+    validator.DeclInputVar("elementwise_" + type + "-Y",
+                           nvinfer1::Dims3(10, 3, 3));
+    validator.DeclOutputVar("elementwise_" + type + "-Out",
+                            nvinfer1::DimsCHW(10, 3, 3));
+
+    // Prepare Op description
+    framework::OpDesc desc;
+    desc.SetType("elementwise_" + type);
+    desc.SetInput("X", {"elementwise_" + type + "-X"});
+    desc.SetInput("Y", {"elementwise_" + type + "-Y"});
+    desc.SetOutput("Out", {"elementwise_" + type + "-Out"});
+
+    int axis = -1;
+    desc.SetAttr("axis", axis);
+
+    validator.SetOp(*desc.Proto());
+    validator.Execute(batch_size);
+  }
+}
 
-  validator.Execute(8);
+TEST(elementwise_op, plugin) {
+  for (std::string type : {"add", "mul"}) {
+    int batch_size = 8;
+    std::unordered_set<std::string> parameters;
+    framework::Scope scope;
+    TRTConvertValidation validator(batch_size, parameters, scope, 1 << 15);
+    validator.DeclInputVar("elementwise_" + type + "-X",
+                           nvinfer1::DimsCHW(10, 3, 3));
+    validator.DeclInputVar("elementwise_" + type + "-Y",
+                           nvinfer1::Dims3(10, 1, 1));
+    validator.DeclOutputVar("elementwise_" + type + "-Out",
+                            nvinfer1::DimsCHW(10, 3, 3));
+
+    // Prepare Op description
+    framework::OpDesc desc;
+    desc.SetType("elementwise_" + type);
+    desc.SetInput("X", {"elementwise_" + type + "-X"});
+    desc.SetInput("Y", {"elementwise_" + type + "-Y"});
+    desc.SetOutput("Out", {"elementwise_" + type + "-Out"});
+
+    int axis = -1;
+    desc.SetAttr("axis", axis);
+
+    validator.SetOp(*desc.Proto());
+    validator.Execute(batch_size);
+  }
 }
 
 }  // namespace tensorrt
 }  // namespace inference
 }  // namespace paddle
+
 USE_OP(elementwise_add);
+USE_OP(elementwise_mul);

paddle/fluid/inference/tensorrt/convert/test_mul_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
+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
+    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. */
+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/framework/op_registry.h"

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

@@ -4,7 +4,7 @@ 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
+    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,

paddle/fluid/inference/tensorrt/engine.cc

@@ -257,9 +257,10 @@ void TensorRTEngine::freshDeviceId() {
 }
 
 nvinfer1::IPluginLayer *TensorRTEngine::AddPlugin(
-    nvinfer1::ITensor *const *inputs, int nbInputs, PluginTensorRT *plugin) {
+    nvinfer1::ITensor *const *inputs, int num_inputs,
+    plugin::PluginTensorRT *plugin) {
   owned_plugin_.emplace_back(plugin);
-  return infer_network_.get()->addPluginExt(inputs, nbInputs, *plugin);
+  return infer_network_.get()->addPluginExt(inputs, num_inputs, *plugin);
 }
 
 }  // namespace tensorrt

paddle/fluid/inference/tensorrt/engine.h

@@ -128,7 +128,7 @@ class TensorRTEngine : public EngineBase {
   int GetRuntimeBatch();
   int GetDevice() { return device_; }
   nvinfer1::IPluginLayer* AddPlugin(nvinfer1::ITensor* const* inputs,
-                                    int nbInputs, PluginTensorRT*);
+                                    int num_inputs, plugin::PluginTensorRT*);
 
   // A pointer to CPU memory is needed of the TRT weight.
   // Before TRT runs, fluid loads weight into GPU storage.
@@ -171,7 +171,7 @@ class TensorRTEngine : public EngineBase {
 
   // The specific GPU id that the TensorRTEngine bounded to.
   int device_;
-  std::vector<std::unique_ptr<PluginTensorRT>> owned_plugin_;
+  std::vector<std::unique_ptr<plugin::PluginTensorRT>> owned_plugin_;
 
   // TensorRT related internal members
   template <typename T>

paddle/fluid/inference/tensorrt/plugin/CMakeLists.txt

-nv_library(tensorrt_plugin SRCS trt_plugin.cc split_op_plugin.cu prelu_op_plugin.cu DEPS enforce)
+nv_library(tensorrt_plugin
+           SRCS trt_plugin.cc split_op_plugin.cu elementwise_op_plugin.cu prelu_op_plugin.cu
+           DEPS enforce device_context)

paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.cu

+/* 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 "paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+namespace details {
+
+template <typename T>
+struct Add {
+  __device__ T operator()(const T& a, const T& b) const { return a + b; }
+};
+
+template <typename T>
+struct Mul {
+  __device__ T operator()(const T& a, const T& b) const { return a * b; }
+};
+
+template <typename T, typename Operator>
+__global__ void ColumnWiseKernel(Operator op, const T* x, const T* y, T* out,
+                                 int batch_size, int num_rows, int num_cols) {
+  for (int batch_id = 0; batch_id < batch_size; ++batch_id) {
+    int row = blockIdx.x;
+    for (; row < num_rows; row += gridDim.x) {
+      T value_y = y[batch_id * num_rows + row];
+      int col = threadIdx.x;
+      int offset = (batch_id * num_rows + row) * num_cols;
+      for (; col < num_cols; col += blockDim.x) {
+        T value_x = x[offset + col];
+        out[offset + col] = op(value_x, value_y);
+      }
+    }
+  }
+}
+
+template <typename T, typename Operator>
+static void ElementWise(Operator op, const T* x, const T* y, T* out,
+                        int batch_size, int prev, int midd, int post,
+                        cudaStream_t stream) {
+  const int kThreadsPerBlock = 1024;
+  const int kMaximumBlocks = 65535;
+  if (prev == 1) {
+    int num_threads = (post > kThreadsPerBlock) ? kThreadsPerBlock
+                                                : (((post + 31) >> 5) << 5);
+    int num_blocks = (midd < kMaximumBlocks) ? midd : kMaximumBlocks;
+    ColumnWiseKernel<<<num_blocks, num_threads, 0, stream>>>(
+        op, x, y, out, batch_size, midd, post);
+  } else if (post == 1) {
+    PADDLE_THROW("Not implemented.");
+  } else {
+    PADDLE_THROW("Not implemented.");
+  }
+}
+
+}  // namespace details
+
+nvinfer1::Dims ElementWisePlugin::getOutputDimensions(
+    int index, const nvinfer1::Dims* input_dims, int num_inputs) {
+  PADDLE_ENFORCE_EQ(index, 0);
+  PADDLE_ENFORCE_EQ(num_inputs, 2);
+  PADDLE_ENFORCE_NOT_NULL(input_dims);
+  return input_dims[0];
+}
+
+int ElementWisePlugin::initialize() {
+  PADDLE_ENFORCE_GT(dims_y_.nbDims, 0);
+
+  axis_ = (axis_ == -1) ? dims_x_.nbDims - dims_y_.nbDims : axis_;
+  int trimed_nb_dims = dims_y_.nbDims;
+  for (; trimed_nb_dims > 0; --trimed_nb_dims) {
+    if (dims_y_.d[trimed_nb_dims - 1] != 1) {
+      break;
+    }
+  }
+  dims_y_.nbDims = trimed_nb_dims;
+
+  PADDLE_ENFORCE_GE(dims_x_.nbDims, dims_y_.nbDims + axis_);
+  PADDLE_ENFORCE_LT(axis_, dims_x_.nbDims);
+
+  prev_size_ = 1;
+  midd_size_ = 1;
+  post_size_ = 1;
+  for (int i = 0; i < axis_; ++i) {
+    prev_size_ *= dims_x_.d[i];
+  }
+
+  for (int i = 0; i < dims_y_.nbDims; ++i) {
+    PADDLE_ENFORCE_EQ(dims_x_.d[i + axis_], dims_y_.d[i],
+                      "Broadcast dimension mismatch.");
+    midd_size_ *= dims_y_.d[i];
+  }
+
+  for (int i = axis_ + dims_y_.nbDims; i < dims_x_.nbDims; ++i) {
+    post_size_ *= dims_x_.d[i];
+  }
+  return 0;
+}
+
+int ElementWisePlugin::enqueue(int batch_size, const void* const* inputs,
+                               void** outputs, void* workspace,
+                               cudaStream_t stream) {
+  const float* x = reinterpret_cast<const float*>(inputs[0]);
+  const float* y = reinterpret_cast<const float*>(inputs[1]);
+  float* out = reinterpret_cast<float*>(outputs[0]);
+
+  if (type_ == nvinfer1::ElementWiseOperation::kSUM) {
+    details::ElementWise(details::Add<float>(), x, y, out, batch_size,
+                         prev_size_, midd_size_, post_size_, stream);
+  } else if (type_ == nvinfer1::ElementWiseOperation::kPROD) {
+    details::ElementWise(details::Mul<float>(), x, y, out, batch_size,
+                         prev_size_, midd_size_, post_size_, stream);
+  } else {
+    PADDLE_THROW("Not implemented.");
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.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 <vector>
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class ElementWisePlugin : public PluginTensorRT {
+ public:
+  ElementWisePlugin(nvinfer1::ElementWiseOperation type,
+                    nvinfer1::Dims const &dims_x, nvinfer1::Dims const &dims_y,
+                    int axis)
+      : type_(type),
+        dims_x_(dims_x),
+        dims_y_(dims_y),
+        axis_(axis),
+        prev_size_(1),
+        midd_size_(1),
+        post_size_(1) {}
+
+  ElementWisePlugin(void const *serial_data, size_t serial_length) {
+    deserializeBase(serial_data, serial_length);
+    DeserializeValue(&serial_data, &serial_length, &axis_);
+    DeserializeValue(&serial_data, &serial_length, &dims_x_);
+    DeserializeValue(&serial_data, &serial_length, &dims_y_);
+  }
+
+  ElementWisePlugin *clone() const override {
+    // return new ElementWisePlugin(dims_x_, dims_y_, axis_);
+    return nullptr;
+  }
+
+  const char *getPluginType() const override { return "elementwise"; }
+
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims *input_dims,
+                                     int num_inputs) override;
+
+  int initialize() override;
+
+  // execute the layer
+  int enqueue(int batch_size, const void *const *inputs, void **outputs,
+              void *workspace, cudaStream_t stream);
+
+ protected:
+  size_t getSerializationSize() override {
+    return SerializedSize(axis_) + SerializedSize(dims_x_) +
+           SerializedSize(dims_y_) + getBaseSerializationSize();
+  }
+
+  void serialize(void *buffer) override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, axis_);
+    SerializeValue(&buffer, dims_x_);
+    SerializeValue(&buffer, dims_y_);
+  }
+
+  nvinfer1::ElementWiseOperation type_;
+  nvinfer1::Dims dims_x_;
+  nvinfer1::Dims dims_y_;
+  int axis_;
+  int prev_size_;
+  int midd_size_;
+  int post_size_;
+};
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/prelu_op_plugin.cu

@@ -20,6 +20,7 @@
 namespace paddle {
 namespace inference {
 namespace tensorrt {
+namespace plugin {
 
 static const int CUDA_NUM_THREADS = 1024;
 static const int CUDA_MAX_NUM_BLOCKS = 65535;
@@ -126,6 +127,7 @@ int PReluPlugin::enqueue(int batchSize, const void *const *inputs,
   return cudaGetLastError() != cudaSuccess;
 }
 
+}  // namespace plugin
 }  // namespace tensorrt
 }  // namespace inference
 }  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/prelu_op_plugin.h

@@ -21,6 +21,7 @@
 namespace paddle {
 namespace inference {
 namespace tensorrt {
+namespace plugin {
 
 class PReluPlugin : public PluginTensorRT {
   TensorRTEngine::Weight alpha_;
@@ -63,6 +64,7 @@ class PReluPlugin : public PluginTensorRT {
               void *workspace, cudaStream_t stream) override;
 };
 
+}  // namespace plugin
 }  // namespace tensorrt
 }  // namespace inference
 }  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/serialize.h

@@ -14,10 +14,15 @@
 
 #pragma once
 
-#include <cassert>
 #include <cstring>
 #include <type_traits>
 #include <vector>
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
 
 template <typename T>
 inline void SerializeValue(void** buffer, T const& value);
@@ -26,7 +31,7 @@ template <typename T>
 inline void DeserializeValue(void const** buffer, size_t* buffer_size,
                              T* value);
 
-namespace {
+namespace details {
 
 template <typename T, class Enable = void>
 struct Serializer {};
@@ -36,10 +41,12 @@ struct Serializer<T, typename std::enable_if<std::is_arithmetic<T>::value ||
                                              std::is_enum<T>::value ||
                                              std::is_pod<T>::value>::type> {
   static size_t SerializedSize(T const& value) { return sizeof(T); }
+
   static void Serialize(void** buffer, T const& value) {
     std::memcpy(*buffer, &value, sizeof(T));
     reinterpret_cast<char*&>(*buffer) += sizeof(T);
   }
+
   static void Deserialize(void const** buffer, size_t* buffer_size, T* value) {
     assert(*buffer_size >= sizeof(T));
     std::memcpy(value, *buffer, sizeof(T));
@@ -51,10 +58,12 @@ struct Serializer<T, typename std::enable_if<std::is_arithmetic<T>::value ||
 template <>
 struct Serializer<const char*> {
   static size_t SerializedSize(const char* value) { return strlen(value) + 1; }
+
   static void Serialize(void** buffer, const char* value) {
-    std::strcpy(static_cast<char*>(*buffer), value);
+    std::strcpy(static_cast<char*>(*buffer), value);  // NOLINT
     reinterpret_cast<char*&>(*buffer) += strlen(value) + 1;
   }
+
   static void Deserialize(void const** buffer, size_t* buffer_size,
                           const char** value) {
     *value = static_cast<char const*>(*buffer);
@@ -73,39 +82,46 @@ struct Serializer<std::vector<T>,
   static size_t SerializedSize(std::vector<T> const& value) {
     return sizeof(value.size()) + value.size() * sizeof(T);
   }
+
   static void Serialize(void** buffer, std::vector<T> const& value) {
     SerializeValue(buffer, value.size());
     size_t nbyte = value.size() * sizeof(T);
     std::memcpy(*buffer, value.data(), nbyte);
     reinterpret_cast<char*&>(*buffer) += nbyte;
   }
+
   static void Deserialize(void const** buffer, size_t* buffer_size,
                           std::vector<T>* value) {
     size_t size;
     DeserializeValue(buffer, buffer_size, &size);
     value->resize(size);
     size_t nbyte = value->size() * sizeof(T);
-    assert(*buffer_size >= nbyte);
+    PADDLE_ENFORCE_GE(*buffer_size, nbyte);
     std::memcpy(value->data(), *buffer, nbyte);
     reinterpret_cast<char const*&>(*buffer) += nbyte;
     *buffer_size -= nbyte;
   }
 };
 
-}  // namespace
+}  // namespace details
 
 template <typename T>
 inline size_t SerializedSize(T const& value) {
-  return Serializer<T>::SerializedSize(value);
+  return details::Serializer<T>::SerializedSize(value);
 }
 
 template <typename T>
 inline void SerializeValue(void** buffer, T const& value) {
-  return Serializer<T>::Serialize(buffer, value);
+  return details::Serializer<T>::Serialize(buffer, value);
 }
 
 template <typename T>
 inline void DeserializeValue(void const** buffer, size_t* buffer_size,
                              T* value) {
-  return Serializer<T>::Deserialize(buffer, buffer_size, value);
+  return details::Serializer<T>::Deserialize(buffer, buffer_size, value);
 }
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/split_op_plugin.cu

@@ -12,26 +12,26 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <stdio.h>
-#include <cassert>
 #include "paddle/fluid/inference/tensorrt/plugin/split_op_plugin.h"
 
 namespace paddle {
 namespace inference {
 namespace tensorrt {
+namespace plugin {
 
-nvinfer1::Dims SplitPlugin::getOutputDimensions(int index,
-                                                const nvinfer1::Dims* inputDims,
-                                                int nbInputs) {
-  assert(nbInputs == 1);
-  assert(index < this->getNbOutputs());
-  nvinfer1::Dims const& input_dims = inputDims[0];
-  nvinfer1::Dims output_dims = input_dims;
+nvinfer1::Dims SplitPlugin::getOutputDimensions(
+    int index, const nvinfer1::Dims* input_dims, int num_inputs) {
+  PADDLE_ENFORCE_EQ(num_inputs, 1);
+  PADDLE_ENFORCE_LT(index, this->getNbOutputs());
+
+  nvinfer1::Dims output_dims = input_dims[0];
   output_dims.d[axis_] = output_length_.at(index);
   return output_dims;
 }
 
 int SplitPlugin::initialize() {
+  PADDLE_ENFORCE_LE(axis_, nvinfer1::Dims::MAX_DIMS);
+
   std::vector<int> segment_offsets(1, 0);
   for (int i = 0; i < this->getNbOutputs(); ++i) {
     segment_offsets.push_back(segment_offsets.back() + output_length_[i]);
@@ -76,6 +76,7 @@ int SplitPlugin::enqueue(int batchSize, const void* const* inputs,
   return cudaGetLastError() != cudaSuccess;
 }
 
-}  // tensorrt
-}  // inference
-}  // paddle
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/split_op_plugin.h

@@ -14,61 +14,58 @@
 
 #pragma once
 
+#include <vector>
 #include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
 
 namespace paddle {
 namespace inference {
 namespace tensorrt {
+namespace plugin {
 
 class SplitPlugin : public PluginTensorRT {
-  int axis_;
-  std::vector<int> output_length_;
-  int nx_, ny_, nz_;
-  std::vector<int> segment_offsets_;
+ public:
+  SplitPlugin(int axis, std::vector<int> const &output_lengths)
+      : axis_(axis), output_length_(output_lengths) {}
+
+  SplitPlugin(void const *serial_data, size_t serial_length) {
+    deserializeBase(serial_data, serial_length);
+    DeserializeValue(&serial_data, &serial_length, &axis_);
+    DeserializeValue(&serial_data, &serial_length, &output_length_);
+  }
+
+  SplitPlugin *clone() const override {
+    return new SplitPlugin(axis_, output_length_);
+  }
+
+  const char *getPluginType() const override { return "split"; }
+  int getNbOutputs() const override { return output_length_.size(); }
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims *input_dims,
+                                     int num_inputs) override;
+
+  int initialize() override;
+  int enqueue(int batchSize, const void *const *inputs, void **outputs,
+              void *workspace, cudaStream_t stream) override;
 
  protected:
-  virtual size_t getSerializationSize() override {
+  size_t getSerializationSize() override {
     return SerializedSize(axis_) + SerializedSize(output_length_) +
            getBaseSerializationSize();
   }
 
-  // TRT will call this func when we need to serialize the configuration of
-  // tensorrt.
-  // It should not be called by users.
-  virtual void serialize(void *buffer) override {
+  void serialize(void *buffer) override {
     serializeBase(buffer);
     SerializeValue(&buffer, axis_);
     SerializeValue(&buffer, output_length_);
   }
 
- public:
-  SplitPlugin(int axis, std::vector<int> const &output_lengths)
-      : axis_(axis), output_length_(output_lengths) {
-    assert(axis <= nvinfer1::Dims::MAX_DIMS);
-  }
-
-  // It was used for tensorrt deserialization.
-  // It should not be called by users.
-  SplitPlugin(void const *serialData, size_t serialLength) {
-    deserializeBase(serialData, serialLength);
-    DeserializeValue(&serialData, &serialLength, &axis_);
-    DeserializeValue(&serialData, &serialLength, &output_length_);
-  }
-
-  SplitPlugin *clone() const override {
-    return new SplitPlugin(axis_, output_length_);
-  }
-
-  virtual const char *getPluginType() const override { return "split"; }
-  virtual int getNbOutputs() const override { return output_length_.size(); }
-  virtual nvinfer1::Dims getOutputDimensions(int index,
-                                             const nvinfer1::Dims *inputs,
-                                             int nbInputDims) override;
-  virtual int initialize() override;
-  virtual int enqueue(int batchSize, const void *const *inputs, void **outputs,
-                      void *workspace, cudaStream_t stream) override;
+  int axis_;
+  std::vector<int> output_length_;
+  int nx_, ny_, nz_;
+  std::vector<int> segment_offsets_;
 };
 
-}  // tensorrt
-}  // inference
-}  // paddle
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/trt_plugin.cc

@@ -17,6 +17,7 @@
 namespace paddle {
 namespace inference {
 namespace tensorrt {
+namespace plugin {
 
 void PluginTensorRT::serializeBase(void*& buffer) {
   SerializeValue(&buffer, input_dims_);
@@ -25,12 +26,12 @@ void PluginTensorRT::serializeBase(void*& buffer) {
   SerializeValue(&buffer, data_format_);
 }
 
-void PluginTensorRT::deserializeBase(void const*& serialData,
-                                     size_t& serialLength) {
-  DeserializeValue(&serialData, &serialLength, &input_dims_);
-  DeserializeValue(&serialData, &serialLength, &max_batch_size_);
-  DeserializeValue(&serialData, &serialLength, &data_type_);
-  DeserializeValue(&serialData, &serialLength, &data_format_);
+void PluginTensorRT::deserializeBase(void const*& serial_data,
+                                     size_t& serial_length) {
+  DeserializeValue(&serial_data, &serial_length, &input_dims_);
+  DeserializeValue(&serial_data, &serial_length, &max_batch_size_);
+  DeserializeValue(&serial_data, &serial_length, &data_type_);
+  DeserializeValue(&serial_data, &serial_length, &data_format_);
 }
 
 size_t PluginTensorRT::getBaseSerializationSize() {
@@ -44,18 +45,17 @@ bool PluginTensorRT::supportsFormat(nvinfer1::DataType type,
           (format == nvinfer1::PluginFormat::kNCHW));
 }
 
-void PluginTensorRT::configureWithFormat(const nvinfer1::Dims* inputDims,
-                                         int nbInputs,
-                                         const nvinfer1::Dims* outputDims,
-                                         int nbOutputs, nvinfer1::DataType type,
-                                         nvinfer1::PluginFormat format,
-                                         int maxBatchSize) {
+void PluginTensorRT::configureWithFormat(
+    const nvinfer1::Dims* input_dims, int num_inputs,
+    const nvinfer1::Dims* output_dims, int num_outputs, nvinfer1::DataType type,
+    nvinfer1::PluginFormat format, int max_batch_size) {
   data_type_ = type;
   data_format_ = format;
-  input_dims_.assign(inputDims, inputDims + nbInputs);
-  max_batch_size_ = maxBatchSize;
+  input_dims_.assign(input_dims, input_dims + num_inputs);
+  max_batch_size_ = max_batch_size;
 }
 
+}  // namespace plugin
 }  // namespace tensorrt
 }  // namespace inference
 }  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/trt_plugin.h

@@ -14,23 +14,30 @@
 
 #pragma once
 
-#include <cassert>
+#include <NvInfer.h>
 #include <cstring>
-#include <iostream>
 #include <unordered_map>
 #include <vector>
-#include "NvInfer.h"
 
 #include "paddle/fluid/inference/tensorrt/plugin/serialize.h"
+#include "paddle/fluid/platform/enforce.h"
+#include "paddle/fluid/platform/profiler.h"
+
+DECLARE_bool(profile);
 
 namespace paddle {
 namespace inference {
 namespace tensorrt {
+namespace plugin {
 
 class PluginTensorRT : public nvinfer1::IPluginExt {
  public:
   PluginTensorRT() {}
+  // It was used for TensorRT deserialization.
+  // It should not be called by users.
   PluginTensorRT(const void* serialized_data, size_t length) {}
+  virtual ~PluginTensorRT() {}
+
   nvinfer1::Dims const& getInputDims(int index) const {
     return input_dims_.at(index);
   }
@@ -38,43 +45,66 @@ class PluginTensorRT : public nvinfer1::IPluginExt {
   nvinfer1::DataType getDataType() const { return data_type_; }
   nvinfer1::PluginFormat getDataFormat() const { return data_format_; }
   virtual const char* getPluginVersion() const { return "1"; }
+
+  void AddInput(nvinfer1::ITensor* input) { inputs_.push_back(input); }
+  std::vector<nvinfer1::ITensor*>& GetInputs() { return inputs_; }
+
+  virtual nvinfer1::IPluginExt* clone() const = 0;
+  virtual const char* getPluginType() const = 0;
+
+  // Following functions are inherit from nvinfer1::IPluginExt
+  // Get the number of outputs from the layer
+  int getNbOutputs() const { return 1; }
+  // Get the dimension of an output tensor
+  virtual nvinfer1::Dims getOutputDimensions(int index,
+                                             const nvinfer1::Dims* input_dims,
+                                             int num_inputs) = 0;
+  // Find the workspace size required by the layer
   size_t getWorkspaceSize(int) const override { return 0; }
+
+  // Initialize the layer for execution.
+  // This is called when the engine is created.
+  int initialize() override { return 0; }
+  // Shutdown the layer. This is called when the engine is destroyed
   void terminate() override {}
-  virtual ~PluginTensorRT() {}
+  // Execute the layer
+  virtual int enqueue(int batch_size, const void* const* inputs, void** outputs,
+                      void* workspace, cudaStream_t stream) = 0;
+
+  // Find the size of the serialization buffer required
+  virtual size_t getSerializationSize() = 0;
+  // Serialize the layer config to buffer.
+  // TensorRT will call this func to serialize the configuration of TensorRT
+  // engine. It should not be called by users.
+  virtual void serialize(void* buffer) = 0;
+
   // Check format support. The default is FLOAT32 and NCHW.
   bool supportsFormat(nvinfer1::DataType type,
                       nvinfer1::PluginFormat format) const override;
-  void configureWithFormat(const nvinfer1::Dims* inputDims, int nbInputs,
-                           const nvinfer1::Dims* outputDims, int nbOutputs,
+  // Configure the layer
+  void configureWithFormat(const nvinfer1::Dims* input_dims, int num_inputs,
+                           const nvinfer1::Dims* output_dims, int num_outputs,
                            nvinfer1::DataType type,
                            nvinfer1::PluginFormat format,
-                           int maxBatchSize) override;
-
-  // *NOTE* The following functions need to be overrided in the subclass.
-  virtual nvinfer1::IPluginExt* clone() const = 0;
-  virtual const char* getPluginType() const = 0;
-  // Initialize the layer for execution. This is called when the engine is
-  // created.
-  int initialize() override { return 0; }
-  // Serialize the layer config to buffer.
-  virtual void serialize(void* buffer) = 0;
-  virtual size_t getSerializationSize() = 0;
-  virtual int enqueue(int batchSize, const void* const* inputs, void** outputs,
-                      void* workspace, cudaStream_t stream) = 0;
+                           int max_batch_size) override;
 
  protected:
   // Deserialize input_dims, max_batch_size, data_type, data_format
-  void deserializeBase(void const*& serialData, size_t& serialLength);
+  void deserializeBase(void const*& serial_data,  // NOLINT
+                       size_t& serial_length);    // NOLINT
   size_t getBaseSerializationSize();
   // Serialize input_dims, max_batch_size, data_type, data_format
-  void serializeBase(void*& buffer);
+  void serializeBase(void*& buffer);  // NOLINT
 
   std::vector<nvinfer1::Dims> input_dims_;
   size_t max_batch_size_;
   nvinfer1::DataType data_type_;
   nvinfer1::PluginFormat data_format_;
+
+  std::vector<nvinfer1::ITensor*> inputs_;
 };
 
+}  // namespace plugin
 }  // namespace tensorrt
 }  // namespace inference
 }  // namespace paddle

paddle/fluid/inference/tests/api/tester_helper.h

@@ -51,7 +51,7 @@ void PrintConfig(const PaddlePredictor::Config *config, bool use_analysis) {
     LOG(INFO) << *reinterpret_cast<const contrib::AnalysisConfig *>(config);
     return;
   }
-  LOG(INFO) << *config;
+  LOG(INFO) << *reinterpret_cast<const NativeConfig *>(config);
 }
 
 void CompareResult(const std::vector<PaddleTensor> &outputs,

paddle/fluid/memory/CMakeLists.txt

 add_subdirectory(detail)
-
-cc_library(malloc SRCS malloc.cc DEPS buddy_allocator place enforce)
+add_subdirectory(allocation)
+cc_library(malloc SRCS malloc.cc DEPS place enforce allocator_facade)
 cc_library(memcpy SRCS memcpy.cc DEPS place)
 
 cc_library(memory
         DEPS
         malloc
         memcpy)
-
-cc_test(malloc_test SRCS malloc_test.cc DEPS malloc)
-
 #if (WITH_GPU)
 #   nv_test(pinned_memory_test SRCS pinned_memory_test.cu  DEPS place memory)
 #endif()

paddle/fluid/memory/allocation/CMakeLists.txt

+cc_library(allocator SRCS allocator.cc DEPS place)
+cc_library(cpu_allocator SRCS cpu_allocator.cc DEPS allocator)
+cc_library(best_fit_allocator SRCS best_fit_allocator.cc DEPS allocator)
+cc_library(locked_allocator SRCS locked_allocator.cc DEPS allocator)
+cc_library(buffered_allocator SRCS buffered_allocator.cc DEPS allocator)
+cc_library(legacy_allocator SRCS legacy_allocator.cc DEPS allocator buddy_allocator)
+cc_test(buffered_allocator_test SRCS buffered_allocator_test.cc DEPS best_fit_allocator locked_allocator buffered_allocator cpu_allocator)
+
+if (WITH_GPU)
+  nv_library(cuda_allocator SRCS cuda_allocator.cc DEPS allocator cuda_device_guard)
+endif()
+
+cc_library(retry_allocator SRCS retry_allocator.cc DEPS allocator)
+
+if (WITH_GPU)
+    nv_test(best_fit_allocator_test
+            SRCS best_fit_allocator_test.cc
+                 best_fit_allocator_test.cu
+            DEPS best_fit_allocator
+                 locked_allocator
+                 cpu_allocator
+                 cuda_allocator
+                 device_context
+                 memcpy)
+else()
+    cc_test(best_fit_allocator_test
+            SRCS best_fit_allocator_test.cc
+            DEPS best_fit_allocator
+                 locked_allocator
+                 cpu_allocator)
+endif()
+
+nv_library(pinned_allocator SRCS pinned_allocator.cc DEPS allocator)
+if (WITH_GPU)
+    set(AllocatorFacadeDeps gpu_info cuda_allocator pinned_allocator cuda_device_guard)
+else ()
+    set(AllocatorFacadeDeps)
+endif()
+
+cc_library(aligned_allocator SRCS aligned_allocator.cc DEPS allocator)
+cc_library(auto_increment_allocator SRCS auto_increment_allocator.cc DEPS allocator)
+cc_library(zero_size_allocator SRCS zero_size_allocator.cc DEPS allocator)
+cc_library(conditional_allocator SRCS conditional_allocator.cc DEPS allocator)
+cc_library(allocator_strategy SRCS allocator_strategy.cc DEPS gflags)
+cc_library(allocator_facade SRCS allocator_facade.cc DEPS
+        ${AllocatorFacadeDeps}
+        cpu_allocator
+        locked_allocator
+        best_fit_allocator
+        aligned_allocator
+        auto_increment_allocator
+        zero_size_allocator
+        conditional_allocator
+        retry_allocator
+        buffered_allocator
+        allocator_strategy
+        legacy_allocator
+        )
+
+nv_test(allocation_and_eigen_test SRCS allocation_and_eigen_test.cu DEPS allocator_facade)
+
+cc_test(retry_allocator_test SRCS retry_allocator_test.cc DEPS retry_allocator best_fit_allocator locked_allocator cpu_allocator)
+
+cc_test(allocator_facade_test SRCS allocator_facade_test.cc DEPS allocator_facade)

paddle/fluid/memory/allocation/aligned_allocator.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/memory/allocation/aligned_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+ThinAlignedAllocator::ThinAlignedAllocator(
+    std::shared_ptr<Allocator> underlyning_allocator)
+    : underlying_allocator_(std::move(underlyning_allocator)) {}
+
+bool ThinAlignedAllocator::IsAllocThreadSafe() const {
+  return underlying_allocator_->IsAllocThreadSafe();
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/aligned_allocator.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 "paddle/fluid/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// The aligned allocation and allocator will wrap a managed allocator,
+// and returns the aligned pointer.
+//
+// NOTE(yy): For speed reason, I just use a template parameter to get
+// alignment, however, it can be an private member if necessary.
+//
+// NOTE(yy): kAlignment must be 2^N. a `static_assert` should be added.
+template <size_t kAlignment>
+class AlignedAllocation : public Allocation {
+  static_assert(kAlignment > 0 && (kAlignment & (kAlignment - 1)) == 0,
+                "kAlignment must be 2^N");
+
+ public:
+  AlignedAllocation(AllocationPtr&& underlying_allocation, size_t size)
+      : Allocation(AlignedPtr(underlying_allocation->ptr()),
+                   size + kAlignment - Offset(underlying_allocation->ptr()),
+                   underlying_allocation->place()),
+        underlying_allocation_(std::move(underlying_allocation)) {}
+
+ private:
+  static void* AlignedPtr(void* ptr) {
+    return reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(ptr) +
+                                   Offset(ptr));
+  }
+
+  // Offset to aligned pointer.
+  // if ptr is already aligned, returns 0.
+  static size_t Offset(void* ptr) {
+    auto ptr_addr = reinterpret_cast<intptr_t>(ptr);
+    intptr_t aligned_addr = (ptr_addr & ~(kAlignment - 1));
+    intptr_t diff = aligned_addr - ptr_addr;
+    if (diff == 0) {
+      return 0;
+    } else {
+      return kAlignment + diff;
+    }
+  }
+
+  AllocationPtr underlying_allocation_;
+};
+
+// Thin aligned allocator is trivial and used to generate a small size binary.
+//
+// NOTE(yy): This is a trick to make a template class. This class extract the
+// common code into a `thin` class. So if there are multiple specification of
+// the template class, the binary size will not extended too much.
+//
+// NOTE(yy): This could be an over design. If it harms readability of code, it
+// could be removed later.
+class ThinAlignedAllocator : public Allocator {
+ public:
+  explicit ThinAlignedAllocator(
+      std::shared_ptr<Allocator> underlyning_allocator);
+
+  bool IsAllocThreadSafe() const;
+
+ protected:
+  std::shared_ptr<Allocator> underlying_allocator_;
+};
+
+// An aligned allocator will allocate `size+kAlignment` allocation and adjust
+// the pointer offset.
+template <size_t kAlignment>
+class AlignedAllocator : public ThinAlignedAllocator {
+ public:
+  using ThinAlignedAllocator::ThinAlignedAllocator;
+
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override {
+    auto raw_allocation =
+        underlying_allocator_->Allocate(size + kAlignment, attr);
+    return new AlignedAllocation<kAlignment>(std::move(raw_allocation), size);
+  }
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocation_and_eigen_test.cu

+// 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/framework/eigen.h"
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/platform/device_context.h"
+#include "paddle/fluid/platform/for_range.h"
+#include "unsupported/Eigen/CXX11/Tensor"
+
+// NOTE(yy): this unittest is not important. It just used for debugging.
+// It can be removed later.
+struct FillZero {
+ public:
+  float* ptr_;
+
+  __device__ void operator()(size_t i) { ptr_[i] = 0.0f; }
+};
+
+namespace paddle {
+TEST(Eigen, main) {
+  framework::Tensor tensor;
+  platform::CUDAPlace gpu(0);
+  float* ptr = tensor.mutable_data<float>({10, 10}, gpu);
+  auto& dev_ctx = *reinterpret_cast<platform::CUDADeviceContext*>(
+      platform::DeviceContextPool::Instance().Get(gpu));
+  PADDLE_ENFORCE(cudaMemset(ptr, 0, sizeof(float) * 100));
+
+  platform::ForRange<platform::CUDADeviceContext> for_range(dev_ctx, 100);
+  for_range(FillZero{ptr});
+  dev_ctx.Wait();
+
+  auto eigen_vec = framework::EigenVector<float>::Flatten(tensor);
+  auto& eigen_dev = *dev_ctx.eigen_device();
+  eigen_vec.device(eigen_dev) = eigen_vec.constant(0.0f);
+}
+}  // namespace paddle

paddle/fluid/memory/allocation/allocation_with_underlying.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/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class AllocationWithUnderlying : public Allocation {
+ public:
+  explicit AllocationWithUnderlying(AllocationPtr allocation)
+      : Allocation(allocation->ptr(), allocation->size(), allocation->place()),
+        allocation_(std::move(allocation)) {}
+  AllocationPtr allocation_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocator.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/memory/allocation/allocator.h"
+
+#include <functional>
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+Allocation::~Allocation() {}
+
+Allocator::~Allocator() {}
+
+bool Allocator::IsAllocThreadSafe() const { return false; }
+
+AllocationPtr Allocator::Allocate(size_t size, Allocator::Attr attr) {
+  auto ptr = AllocateImpl(size, attr);
+  ptr->set_allocator(this);
+  return AllocationPtr(ptr);
+}
+
+void Allocator::Free(Allocation* allocation) { delete allocation; }
+
+const char* BadAlloc::what() const noexcept { return msg_.c_str(); }
+
+void AllocationDeleter::operator()(Allocation* allocation) const {
+  auto* allocator = allocation->allocator();
+  allocator->Free(allocation);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocator.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 "paddle/fluid/platform/place.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// Exception when `Alloc`/`AllocShared` failed
+class BadAlloc : public std::exception {
+ public:
+  explicit BadAlloc(std::string msg) : msg_(std::move(msg)) {}
+  const char* what() const noexcept override;
+
+ private:
+  std::string msg_;
+};
+
+class Allocation;
+class AllocationDeleter {
+ public:
+  void operator()(Allocation* allocation) const;
+};
+
+class Allocator;
+// Allocation is the object holding the actually pointer. Use
+// `Allocation::ptr()` will returns the pointer that allocated.
+//
+// NOTE: this is the base class of Allocation. Each allocator can use its own
+//       allocation object.
+// NOTE: the `Allocation::ptr()` could be nullptr, if the allocation size is 0
+class Allocation {
+ public:
+  Allocation(void* ptr, size_t size, platform::Place place)
+      : allocator_(nullptr), ptr_(ptr), size_(size), place_(place) {}
+
+  Allocation(const Allocation& o) = delete;
+  Allocation& operator=(const Allocation& o) = delete;
+
+  // Returns the holding pointer.
+  // NOTE: For performance consideration, it is better not to make this method
+  // as a virtual method. If we want to implement a `defragmentation` later,
+  // we might need to make `ptr_` field as a protected field, and add a virtual
+  // method like `defragmentation` to change `ptr_`.
+  void* ptr() const { return ptr_; }
+
+  // Returns the size of this memory buffer, i.e., ptr() + size() - 1 is the
+  // last valid element.
+  //
+  // NOTE: Some allocator might alloc more memory than request. The size
+  // could larger than its request. For example,
+  //    the AlignedAllocator will always allocate memory as size + kAlignment.
+  //    The raw pointer might not aligned, so an offset might be added to raw
+  //    the pointer. The size of this allocation will be
+  //    `size + kAlignemnt - offset`.
+  size_t size() const { return size_; }
+
+  const platform::Place& place() const { return place_; }
+
+  Allocator* allocator() { return allocator_; }
+
+  void set_allocator(Allocator* allocator) { allocator_ = allocator; }
+
+  virtual ~Allocation();
+
+ private:
+  Allocator* allocator_;
+  void* ptr_;
+  size_t size_;
+  platform::Place place_;
+};
+
+using AllocationPtr = std::unique_ptr<Allocation, AllocationDeleter>;
+
+// Base interface class of memory Allocator.
+// To allocate a memory, allocator needs two parameters:
+//    1. size of bytes.
+//    2. Attribute of memory.
+// NOTE: the attribute of memory might be ignored if the allocator does not
+// care it.
+class Allocator {
+ public:
+  enum Attr {
+    kDefault = 0,  // Default attribute. Uses the fast or stablest allocation
+                   // algorithm.
+
+    kFixedHuge = 1,  // The allocation may not be freed until the program
+                     // ends. e.g., `Parameters` and `Momentum`.
+
+    kFluxHuge = 2,  // The allocation may create and freed frequently and the
+                    // allocation is considerable huge. Like `activations`
+                    // and gradients.
+
+    kScratchpad =
+        3,  // The `Scratchpad` memory is allocated and freed very soon,
+            // usually within an operator or aux memory.
+            // Like CUDNN workspace, AUX memory in batch norm, etc.
+            //
+            // https://en.wikipedia.org/wiki/Scratchpad_memory
+
+    kCrossDevice =
+        4,  // The memory used cross-device memory copy/communication.
+            // For example:
+            // 1. it can use an `pinned` memory for CPU-GPU
+            //    communication.
+            // 2. it can use an `registered` memory for RDMA
+            //    communication.
+
+    NumOfAttrs = 5  // The number of all attributes. It is used internally.
+  };
+
+  virtual ~Allocator();
+
+  // Allocate an allocation.
+  AllocationPtr Allocate(size_t size, Allocator::Attr attr = kDefault);
+
+  // True if the `Allocate` is thread safe.
+  virtual bool IsAllocThreadSafe() const;
+
+ protected:
+  virtual void Free(Allocation* allocation);
+  virtual Allocation* AllocateImpl(size_t size, Allocator::Attr attr) = 0;
+
+ private:
+  friend class AllocationDeleter;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocator_facade.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/memory/allocation/allocator.h"
+#include <gflags/gflags.h>
+#include <map>
+#include <string>
+#include <unordered_map>
+#include <vector>
+#include "paddle/fluid/memory/allocation/aligned_allocator.h"
+#include "paddle/fluid/memory/allocation/allocator_facade.h"
+#include "paddle/fluid/memory/allocation/allocator_strategy.h"
+#include "paddle/fluid/memory/allocation/auto_increment_allocator.h"
+#include "paddle/fluid/memory/allocation/best_fit_allocator.h"
+#include "paddle/fluid/memory/allocation/conditional_allocator.h"
+#include "paddle/fluid/memory/allocation/cpu_allocator.h"
+#include "paddle/fluid/memory/allocation/legacy_allocator.h"
+#include "paddle/fluid/memory/allocation/locked_allocator.h"
+#include "paddle/fluid/memory/allocation/retry_allocator.h"
+#include "paddle/fluid/memory/allocation/zero_size_allocator.h"
+#include "paddle/fluid/platform/cpu_info.h"
+#include "paddle/fluid/platform/place.h"
+#ifdef PADDLE_WITH_CUDA
+#include "paddle/fluid/memory/allocation/cuda_allocator.h"
+#include "paddle/fluid/memory/allocation/pinned_allocator.h"
+#include "paddle/fluid/platform/cuda_device_guard.h"
+#include "paddle/fluid/platform/gpu_info.h"
+#endif
+
+DEFINE_int64(
+    gpu_allocator_retry_time, 0,
+    "The retry time (milliseconds) when allocator fails "
+    "to allocate memory. No retry if this value is not greater than 0");
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// TODO(yy): Dirty code here. This class should be configurable in runtime.
+class CPUManagedAllocator : public Allocator {
+ public:
+  CPUManagedAllocator() : normal_allocator_(new CPUAllocator()) {}
+
+  bool IsAllocThreadSafe() const override { return true; }
+
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override {
+    return normal_allocator_->Allocate(size, attr).release();
+  }
+
+ private:
+  std::shared_ptr<Allocator> normal_allocator_;
+};
+
+// TODO(yy): Dirty code here. This class should be configurable in runtime.
+class ChunkedAllocator : public Allocator {
+ public:
+  explicit ChunkedAllocator(std::unique_ptr<Allocator> system_allocator,
+                            size_t max_chunk_size, size_t capacity = 1,
+                            int64_t retry_time = -1)
+      : max_chunk_size_(max_chunk_size), retry_time_(retry_time) {
+    raw_allocator_ = std::move(system_allocator);
+
+    if (max_chunk_size_ == 0) {
+      default_allocator_ = raw_allocator_;
+    } else {
+      if (capacity == 1) {
+        VLOG(10) << "Create BestFitAllocator with chunk_size "
+                 << max_chunk_size_;
+        default_allocator_ = CreateAllocatorWithChunk();
+      } else {
+        VLOG(10) << "Create AutoIncrementAllocator with chunk_size "
+                 << max_chunk_size_ << " and capacity " << capacity;
+        default_allocator_ = std::make_shared<AutoIncrementAllocator>(
+            [this] { return std::move(CreateAllocatorWithChunk()); }, capacity);
+      }
+    }
+
+    auto* cond_allocator = new ConditionalAllocator();
+    cond_allocator
+        ->AddAllocator(
+            [this](size_t size, Attr attr) { return size < max_chunk_size_; },
+            default_allocator_)
+        .AddAllocator(
+            [](size_t size, Attr attr) {
+              return true;  // default case
+            },
+            raw_allocator_);
+    default_allocator_.reset(cond_allocator);
+  }
+
+  ~ChunkedAllocator() override {
+    // Specify destruct order.
+    default_allocator_.reset();
+    chunks_.clear();
+    raw_allocator_.reset();
+  }
+
+  std::shared_ptr<Allocator> CreateAllocatorWithChunk() {
+    chunks_.emplace_back(raw_allocator_->Allocate(max_chunk_size_));
+    auto* allocation = chunks_.back().get();
+    std::unique_ptr<Allocator> allocator(new LockedAllocator(
+        std::unique_ptr<Allocator>(new BestFitAllocator(allocation))));
+
+    if (retry_time_ > 0) {
+      auto* retry_allocator =
+          new RetryAllocator(std::move(allocator), retry_time_);
+      allocator.reset(retry_allocator);
+    }
+
+    return std::make_shared<AlignedAllocator<64u>>(std::move(allocator));
+  }
+
+  bool IsAllocThreadSafe() const override { return true; }
+
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override {
+    return default_allocator_->Allocate(size, attr).release();
+  }
+
+ protected:
+  size_t max_chunk_size_;
+  int64_t retry_time_;
+  std::vector<AllocationPtr> chunks_;
+  std::shared_ptr<Allocator> raw_allocator_;
+  std::shared_ptr<Allocator> default_allocator_;
+};
+
+#ifdef PADDLE_WITH_CUDA
+
+class CUDAChunkedAllocator : public ChunkedAllocator {
+ public:
+  explicit CUDAChunkedAllocator(int dev_id)
+      : ChunkedAllocator(std::unique_ptr<Allocator>(
+                             new CUDAAllocator(platform::CUDAPlace(dev_id))),
+                         GetMaxChunkSize(dev_id), GetCapcity(dev_id),
+                         GetRetryTime()) {}
+
+ private:
+  static size_t GetMaxChunkSize(int dev_id) {
+    platform::CUDADeviceGuard guard(dev_id);
+    return platform::GpuMaxChunkSize();
+  }
+
+  static size_t GetCapcity(int dev_id) {
+    platform::CUDADeviceGuard guard(dev_id);
+    size_t available, total;
+    platform::GpuMemoryUsage(&available, &total);
+    size_t max_chunk_size = platform::GpuMaxChunkSize();
+    return max_chunk_size == 0 ? 0 : available / max_chunk_size;
+  }
+
+  static int64_t GetRetryTime() { return FLAGS_gpu_allocator_retry_time; }
+};
+
+class CUDAPinnedChunkedAllocator : public ChunkedAllocator {
+ public:
+  CUDAPinnedChunkedAllocator()
+      : ChunkedAllocator(std::unique_ptr<Allocator>(new CPUPinnedAllocator()),
+                         platform::CUDAPinnedMaxChunkSize(), GetCapacity(),
+                         -1) {}  // never retry
+
+ private:
+  static size_t GetCapacity() {
+    size_t total = platform::CpuTotalPhysicalMemory();
+    size_t max_chunk_size = platform::CUDAPinnedMaxChunkSize();
+    return max_chunk_size == 0 ? 0 : total / max_chunk_size;
+  }
+};
+
+#endif
+
+class AllocatorFacadePrivate {
+ public:
+  std::map<platform::Place, std::shared_ptr<Allocator>> allocators_;
+
+  ~AllocatorFacadePrivate() = default;
+
+  AllocatorFacadePrivate() {
+    if (GetAllocatorStrategy() == AllocatorStrategy::kLegacy) {
+      InitLegacyAllocator();
+    } else {
+      InitCPUAllocator();
+      InitCUDAAllocator();
+      InitCUDAPinnedAllocator();
+      WrapZeroSizeAllocator();
+    }
+  }
+
+ private:
+  void InitLegacyAllocator() {
+    std::vector<platform::Place> places{platform::CPUPlace()};
+#ifdef PADDLE_WITH_CUDA
+    for (int dev_id = 0; dev_id < platform::GetCUDADeviceCount(); ++dev_id) {
+      places.emplace_back(platform::CUDAPlace(dev_id));
+    }
+    places.emplace_back(platform::CUDAPinnedPlace());
+#endif
+    for (auto& p : places) {
+      allocators_[p] = std::make_shared<LegacyAllocator>(p);
+    }
+  }
+
+  void InitCPUAllocator() {
+    allocators_[platform::CPUPlace()] = std::make_shared<CPUManagedAllocator>();
+  }
+
+  void InitCUDAAllocator() {
+#ifdef PADDLE_WITH_CUDA
+    int device_count = platform::GetCUDADeviceCount();
+    for (int dev_id = 0; dev_id < device_count; ++dev_id) {
+      allocators_[platform::CUDAPlace(dev_id)] =
+          std::make_shared<CUDAChunkedAllocator>(dev_id);
+    }
+#endif
+  }
+
+  void InitCUDAPinnedAllocator() {
+#ifdef PADDLE_WITH_CUDA
+    allocators_[platform::CUDAPinnedPlace()] =
+        std::make_shared<CUDAPinnedChunkedAllocator>();
+#endif
+  }
+
+  void WrapZeroSizeAllocator() {
+    for (auto& pair : allocators_) {
+      pair.second =
+          std::make_shared<ZeroSizeAllocator>(pair.second, pair.first);
+    }
+  }
+};
+
+// Pimpl. Make interface clean.
+AllocatorFacade::AllocatorFacade() : m_(new AllocatorFacadePrivate()) {}
+AllocatorFacade::~AllocatorFacade() { delete m_; }
+
+AllocatorFacade& AllocatorFacade::Instance() {
+  static AllocatorFacade instance;
+  return instance;
+}
+
+std::shared_ptr<Allocation> AllocatorFacade::AllocShared(
+    const platform::Place& place, size_t size, Allocator::Attr attr) {
+  return std::shared_ptr<Allocation>(Alloc(place, size, attr).release(),
+                                     AllocationDeleter());
+}
+
+AllocationPtr AllocatorFacade::Alloc(const platform::Place& place, size_t size,
+                                     Allocator::Attr attr) {
+  auto it = m_->allocators_.find(place);
+  if (it == m_->allocators_.end()) {
+    throw BadAlloc(
+        string::Sprintf("No such allocator for the place, %s", place));
+  }
+  return m_->allocators_.at(place)->Allocate(size, attr);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocator_facade.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 "paddle/fluid/memory/allocation/allocator.h"
+#include "paddle/fluid/platform/place.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// Allocator Facade is the interface exposed to other modules.
+// All the configuration or dirty code under development should
+// be hidden behind this facade.
+//
+// NOTE(yy): This class is a singleton class.
+// NOTE(yy): To create a stable ABI and make compilation faster. Here we use
+// a Pimpl trick;
+class AllocatorFacadePrivate;
+class AllocatorFacade {
+ public:
+  ~AllocatorFacade();
+  AllocatorFacade(const AllocatorFacade& o) = delete;
+  const AllocatorFacade& operator=(const AllocatorFacade& o) = delete;
+
+  static AllocatorFacade& Instance();
+
+  // Allocate a shared allocation.
+  std::shared_ptr<Allocation> AllocShared(
+      const platform::Place& place, size_t size,
+      Allocator::Attr attr = Allocator::kDefault);
+
+  // Allocate a unique allocation.
+  AllocationPtr Alloc(const platform::Place& place, size_t size,
+                      Allocator::Attr attr = Allocator::kDefault);
+
+  // TODO(yy): Allocate a Copy-On-Write allocation?
+ private:
+  AllocatorFacade();
+  AllocatorFacadePrivate* m_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocator_facade_test.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/memory/allocation/allocator_facade.h"
+#include <gflags/gflags.h>
+#include <gtest/gtest.h>
+
+#ifdef PADDLE_WITH_CUDA
+DECLARE_double(fraction_of_gpu_memory_to_use);
+DECLARE_double(fraction_of_cuda_pinned_memory_to_use);
+DECLARE_int64(gpu_allocator_retry_time);
+#endif
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(allocator, allocator) {
+#ifdef PADDLE_WITH_CUDA
+  FLAGS_fraction_of_gpu_memory_to_use = 0.01;
+  FLAGS_gpu_allocator_retry_time = 500;
+  FLAGS_fraction_of_cuda_pinned_memory_to_use = 0.5;
+#endif
+
+  auto &instance = AllocatorFacade::Instance();
+  platform::Place place;
+  size_t size = 1024;
+
+  {
+    place = platform::CPUPlace();
+    size = 1024;
+    auto cpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(cpu_allocation, nullptr);
+    ASSERT_NE(cpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(cpu_allocation->place(), place);
+    ASSERT_EQ(cpu_allocation->size(), size);
+  }
+
+#ifdef PADDLE_WITH_CUDA
+  {
+    place = platform::CUDAPlace(0);
+    size = 1024;
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(), size);
+  }
+
+  {
+    // Allocate 2GB gpu memory
+    place = platform::CUDAPlace(0);
+    size = 2 * static_cast<size_t>(1 << 30);
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(), size);
+  }
+
+  {
+    place = platform::CUDAPinnedPlace();
+    size = (1 << 20);
+    auto cuda_pinned_allocation =
+        instance.Alloc(platform::CUDAPinnedPlace(), 1 << 20);
+    ASSERT_NE(cuda_pinned_allocation, nullptr);
+    ASSERT_NE(cuda_pinned_allocation->ptr(), nullptr);
+    ASSERT_EQ(cuda_pinned_allocation->place(), place);
+    ASSERT_GE(cuda_pinned_allocation->size(), size);
+  }
+#endif
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocator_strategy.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/memory/allocation/allocator_strategy.h"
+#include "gflags/gflags.h"
+
+DEFINE_string(
+    allocator_strategy, "legacy",
+    "The allocation strategy. Legacy means the original allocator of Fluid."
+    "New means the experimental allocators of Fluid. in [legacy, new]");
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+static AllocatorStrategy GetStrategyFromFlag() {
+  return FLAGS_allocator_strategy == "legacy"
+             ? AllocatorStrategy::kLegacy
+             : AllocatorStrategy::kNaiveBestFit;
+}
+
+AllocatorStrategy GetAllocatorStrategy() {
+  static AllocatorStrategy strategy = GetStrategyFromFlag();
+  return strategy;
+}
+
+void UseAllocatorStrategyGFlag() {}
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocator_strategy.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
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+enum class AllocatorStrategy { kLegacy, kNaiveBestFit };
+
+extern AllocatorStrategy GetAllocatorStrategy();
+
+// Do nothing, just make sure linker do not prune this file.
+extern void UseAllocatorStrategyGFlag();
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/auto_increment_allocator.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/memory/allocation/auto_increment_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+bool AutoIncrementAllocator::IsAllocThreadSafe() const { return true; }
+
+std::shared_ptr<Allocator> AutoIncrementAllocator::CreateNewAllocator() {
+  std::lock_guard<std::mutex> guard(mtx_);
+  auto old_size = allocator_num_.load();
+  PADDLE_ENFORCE_LT(old_size, underlying_allocators_.size(),
+                    "Allocator number exceeds capacity %d",
+                    underlying_allocators_.size());
+  underlying_allocators_[old_size] = creator_();
+  prev_success_allocator_ = old_size;
+  ++allocator_num_;
+  PADDLE_ENFORCE(
+      underlying_allocators_[old_size]->IsAllocThreadSafe(),
+      "the underlying allocator must be thread safe. This is a program "
+      "bug.");
+  return underlying_allocators_[old_size];
+}
+Allocation *AutoIncrementAllocator::AllocateImpl(size_t size,
+                                                 Allocator::Attr attr) {
+  auto cur = prev_success_allocator_.load();
+  size_t retry_count = allocator_num_.load();
+  size_t allocator_num = retry_count;
+  while (retry_count-- > 0) {  // until there retry count is zero
+    try {
+      auto res = underlying_allocators_[cur]->Allocate(size, attr);
+      prev_success_allocator_ = cur;
+      return res.release();
+    } catch (BadAlloc &) {
+      if (++cur >= allocator_num) {
+        cur = 0;
+      }
+    } catch (...) {
+      // if there is another type of allocation, just rethrow it.
+      throw;
+    }
+  }
+
+  // This happens when the first allocator is exhausted and
+  // there are more than 1 allocation requests
+  // In this situation, the first allocation request would success
+  // and the second allocation request would fail if we do not use
+  // the newly created allocator by the first allocation request.
+  for (cur = allocator_num; cur < allocator_num_; ++cur) {
+    try {
+      auto ret = underlying_allocators_[cur]->Allocate(size, attr);
+      prev_success_allocator_ = cur;
+      return ret.release();
+    } catch (BadAlloc &) {
+    } catch (...) {
+      throw;
+    }
+  }
+  // No suitable allocator
+  return CreateNewAllocator()->Allocate(size, attr).release();
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/auto_increment_allocator.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 <atomic>  // NOLINT
+#include <functional>
+#include <memory>
+#include <mutex>   // NOLINT
+#include <thread>  // NOLINT
+#include <vector>
+#include "paddle/fluid/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// The AutoIncrementAllocator manages many underlying allocators. If none of
+// them can allocate the request memory, a new allocator will be created and
+// invoke its `allocate` method.
+//
+// NOTE(yy): The AutoIncrementAllocator will prefer to allocate memory from
+// the latest successful allocator.
+//
+// NOTE(yy): We may need to release an underlying allocator if it allocate
+// nothing. However, it is generally not useful, since it will make performance
+// undetermined.
+//
+// NOTE(yy): This allocator is only locked when creating new underlying
+// allocator. The allocation requests from many threads may be dispatched
+// to the same underlying allocator. So the underlying allocator must be
+// thread safe.
+//
+// NOTE(zjl): Add capacity parameters to constructor. A high-performance
+// thread-safe std::vector with varying size is hard to implement.
+// Fortunately, we can get the total GPU memory and each chunk size.
+// Therefore, we can get the suitable capacity of AutoIncrementAllocator.
+class AutoIncrementAllocator : public Allocator {
+ public:
+  // Creator is the method to create ManagedAllocator
+  using AllocatorCreator = std::function<std::shared_ptr<Allocator>()>;
+
+  explicit AutoIncrementAllocator(AllocatorCreator&& creator, size_t capacity)
+      : creator_(std::move(creator)), underlying_allocators_(capacity) {}
+
+  bool IsAllocThreadSafe() const override;
+
+ private:
+  std::shared_ptr<Allocator> CreateNewAllocator();
+
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
+
+ private:
+  AllocatorCreator creator_;
+
+  std::vector<AllocatorCreator::result_type> underlying_allocators_;
+  std::atomic<size_t> allocator_num_{0};
+
+  // Use std::atomic rather than std::mutex, since std::atomic is usually
+  // lock-free
+  std::atomic<size_t> prev_success_allocator_{0};
+
+  std::mutex mtx_;
+};
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/best_fit_allocator.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/memory/allocation/best_fit_allocator.h"
+#include <cmath>
+#include <list>
+#include <map>
+#include <string>
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+static int HighestBitPos(size_t N) {
+  if (UNLIKELY(N == 0)) {
+    return 0;
+  } else {
+#ifdef __GNUCC__
+    return sizeof(unsigned int) * 8 - __builtin_clz(N);
+#else
+    return static_cast<int>(std::log2(N) + 1);
+#endif
+  }
+}
+
+BestFitAllocator::BestFitAllocator(Allocation* allocation)
+    : allocation_(allocation) {
+  details::Chunk chunk;
+  chunk.size_ = allocation_->size();
+  chunk.offset_ = 0;
+  chunk.is_free = true;
+  chunks_.emplace_back(chunk);
+  free_chunks_[HighestBitPos(chunk.size_)].insert(
+      {chunk.size_, chunks_.begin()});
+}
+
+size_t BestFitAllocator::FreeSize() const {
+  size_t acc = 0;
+  for (auto& array_item : free_chunks_) {
+    for (auto& pair : array_item) {
+      acc += pair.second->size_;
+    }
+  }
+  return acc;
+}
+
+BestFitAllocator::ListIt BestFitAllocator::SplitChunk(size_t request_size,
+                                                      size_t free_chunk_offset,
+                                                      MapIt bin_iterator) {
+  auto to_split_it = bin_iterator->second;
+  free_chunks_[free_chunk_offset].erase(bin_iterator);
+
+  PADDLE_ENFORCE(to_split_it->is_free);
+  PADDLE_ENFORCE_GE(to_split_it->size_, request_size);
+
+  auto remaining_size = to_split_it->size_ - request_size;
+  details::Chunk to_use;
+  details::Chunk remaining;
+  to_use.size_ = request_size;
+  to_use.is_free = false;
+  remaining.size_ = remaining_size;
+  remaining.is_free = true;
+
+  // calc offsets
+  to_use.offset_ = to_split_it->offset_;
+  remaining.offset_ = to_use.offset_ + to_use.size_;
+
+  // insert to chunk list
+  auto to_use_it = chunks_.insert(to_split_it, to_use);
+  if (remaining.size_ != 0) {
+    auto bit_size = static_cast<size_t>(HighestBitPos(remaining.size_));
+    free_chunks_[bit_size].insert(
+        {remaining.size_, chunks_.insert(to_split_it, remaining)});
+  }
+  chunks_.erase(to_split_it);
+  return to_use_it;
+}
+
+void BestFitAllocator::InsertFreeNode(const ListIt& it) {
+  auto pos = static_cast<size_t>(HighestBitPos(it->size_));
+  auto& free_map = free_chunks_[pos];
+  free_map.insert({it->size_, it});
+}
+void BestFitAllocator::EraseFreeNode(const ListIt& it) {
+  size_t pos = static_cast<size_t>(HighestBitPos(it->size_));
+  auto& free_map = free_chunks_[pos];
+  auto map_it = free_map.find(it->size_);
+  while (map_it->second != it && map_it != free_map.end()) {
+    ++map_it;
+  }
+  PADDLE_ENFORCE(map_it != free_map.end());
+  free_map.erase(map_it);
+}
+size_t BestFitAllocator::NumFreeChunks() const {
+  size_t num = 0;
+  for (auto& array_item : free_chunks_) {
+    num += array_item.size();
+  }
+  return num;
+}
+void BestFitAllocator::Free(Allocation* allocation) {
+  auto* bf_allocation = dynamic_cast<BestFitAllocation*>(allocation);
+  auto chunk_it = bf_allocation->ChunkIterator();
+  PADDLE_ENFORCE(!chunk_it->is_free);
+  chunk_it->is_free = true;
+  if (chunk_it != chunks_.begin()) {
+    auto prev_it = chunk_it;
+    --prev_it;
+
+    if (prev_it->is_free) {
+      // Merge Left.
+      EraseFreeNode(prev_it);
+      prev_it->size_ += chunk_it->size_;
+      chunks_.erase(chunk_it);
+      chunk_it = prev_it;
+    }
+  }
+
+  auto next_it = chunk_it;
+  ++next_it;
+  if (next_it != chunks_.end() && next_it->is_free) {
+    EraseFreeNode(next_it);
+    chunk_it->size_ += next_it->size_;
+    chunks_.erase(next_it);
+  }
+
+  InsertFreeNode(chunk_it);
+  delete allocation;
+}
+Allocation* BestFitAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
+  auto highest_set_bit = static_cast<size_t>(HighestBitPos(size));
+  MapIt map_it;
+  for (; highest_set_bit < free_chunks_.size(); ++highest_set_bit) {
+    map_it = free_chunks_[highest_set_bit].lower_bound(size);
+    if (map_it != free_chunks_[highest_set_bit].end()) {
+      break;
+    }
+  }
+  if (UNLIKELY(highest_set_bit == free_chunks_.size())) {
+    throw BadAlloc(string::Sprintf(
+        "Cannot allocate %d, All fragments size is %d", size, FreeSize()));
+  }
+  auto chunk_it = SplitChunk(size, highest_set_bit, map_it);
+  return new BestFitAllocation(this, chunk_it);
+}
+
+BestFitAllocation::BestFitAllocation(
+    paddle::memory::allocation::BestFitAllocator* allocator,
+    typename details::ChunkList::iterator chunk_it)
+    : Allocation(reinterpret_cast<void*>(
+                     reinterpret_cast<uintptr_t>(allocator->BasePtr()) +
+                     chunk_it->offset_),
+                 chunk_it->size_, allocator->Place()),
+      chunk_it_(chunk_it) {}
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/best_fit_allocator.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 <array>
+#include <list>
+#include <map>
+#include "paddle/fluid/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+namespace details {
+struct Chunk {
+  bool is_free{true};
+  // Offset to the base allocation.
+  uintptr_t offset_;
+  size_t size_;
+};
+
+// Here we use std::list to maintain chunk list.
+// NOTE(yy): The traditional implementation of ChunkList is add `prev`/`next`
+// pointers in `Chunk`, and split the allocation as `ChunkHeader` and
+// `Payload`. Such as
+//   *-------*---------------*---------------*--------------*
+//   | Chunk | prev_ pointer | next_ pointer | payload .... |
+//   *-------*---------------*---------------*--------------*
+// This implementation can just return a raw pointer, and we can get the list
+// structure by the raw pointer. However, we cannot use the same code on GPU
+// since CPU cannot access GPU memory directly.
+//
+// So we choose to use `std::list` and return an allocation instance, which
+// contains the list node iterator, then we can unify CPU/GPU code.
+//
+// To return an allocation is not a bad idea, since Tensor/Vector should holds
+// an allocation instead of raw pointer directly.
+using ChunkList = std::list<Chunk>;
+
+// Here we use a multi-level map of free chunks.
+// the map is
+//      MSB offset --> size --> [ChunkList::iterator]
+//
+// The time complexities:
+//     find a free chunk:
+//          O(logN),
+//               where N is the number of free nodes with the same MSB offset.
+//     find the position of a chunk iterator:
+//          O(logN + K),
+//               where N is the number of free nodes with the same MSB offset.
+//               where K is the number of free nodes with the same size.
+//     insert a free chunk:
+//          O(logN),
+//               where N is the number of free nodes with the same MSB offset.
+//     erase a free chunk:
+//          O(1)
+using FreeChunkBin =
+    std::array<std::multimap<size_t, ChunkList::iterator>, sizeof(size_t) * 8>;
+}  // namespace details
+
+class BestFitAllocator;
+
+// The BestFitAllocation maintain the List Node iterator.
+class BestFitAllocation : public Allocation {
+ private:
+  using ListIt = typename details::ChunkList::iterator;
+
+ public:
+  BestFitAllocation(BestFitAllocator* allocator, ListIt chunk_it);
+
+  const ListIt& ChunkIterator() const { return chunk_it_; }
+
+ private:
+  typename details::ChunkList::iterator chunk_it_;
+};
+
+// TODO(yy): Current BestFitAllocator is not thread-safe. To make it thread
+// safe, we must wrap a locked_allocator. However, we can implement a thread
+// safe allocator by locking each bin and chunks list independently. It will
+// make BestFitAllocator faster in multi-thread situation.
+//
+// This allocator implements a best-fit allocator with merging the free nodes.
+//
+// To allocate a buffer, it will find the best-fit chunk. If the best-fit chunk
+// is larger than request size, the original block will be split into two
+// chunks. The first block will be used and the second block will be put into
+// free chunks.
+//
+// To free an allocation, it will set the chunk of allocation to free and merge
+// the prev-chunk and the next-chunk when possible.
+class BestFitAllocator : public Allocator {
+ public:
+  explicit BestFitAllocator(Allocation* allocation);
+
+  void* BasePtr() const { return allocation_->ptr(); }
+
+  const platform::Place& Place() const { return allocation_->place(); }
+
+  size_t NumFreeChunks() const;
+
+ private:
+  size_t FreeSize() const;
+  using MapIt = typename details::FreeChunkBin::value_type::iterator;
+  using ListIt = typename details::ChunkList::iterator;
+
+  ListIt SplitChunk(size_t request_size, size_t free_chunk_offset,
+                    MapIt bin_iterator);
+  void EraseFreeNode(const ListIt& it);
+  void InsertFreeNode(const ListIt& it);
+
+ protected:
+  void Free(Allocation* allocation) override;
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
+
+ private:
+  Allocation* allocation_;  // not owned
+  details::ChunkList chunks_;
+  details::FreeChunkBin free_chunks_;
+};
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/best_fit_allocator_test.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/memory/allocation/best_fit_allocator.h"
+#include <thread>  // NOLINT
+#include <vector>
+#include "gtest/gtest.h"
+#include "paddle/fluid/memory/allocation/cpu_allocator.h"
+#include "paddle/fluid/memory/allocation/locked_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class StubAllocation : public Allocation {
+ public:
+  explicit StubAllocation(size_t size)
+      : Allocation(0, size, platform::CPUPlace()) {}
+};
+
+TEST(BestFitAllocator, test_allocation) {
+  StubAllocation stub(4UL * 1024 * 1024 * 1024);
+  BestFitAllocator allocator(&stub);
+  { auto allocation = allocator.Allocate(64, allocator.kDefault); }
+
+  {
+    auto allocation = allocator.Allocate(80, allocator.kDefault);
+
+    {
+      auto best_fit_allocation =
+          dynamic_cast<BestFitAllocation*>(allocation.get());
+      ASSERT_NE(best_fit_allocation, nullptr);
+      ASSERT_FALSE(best_fit_allocation->ChunkIterator()->is_free);
+      ASSERT_EQ(best_fit_allocation->ChunkIterator()->offset_, 0);
+      ASSERT_EQ(allocation->size(), 80);
+      ASSERT_EQ(allocation->ptr(), nullptr);
+    }
+
+    auto allocation2 = allocator.Allocate(60, allocator.kDefault);
+    auto allocation3 = allocator.Allocate(90, allocator.kDefault);
+    allocation2.reset();
+    allocation2 = allocator.Allocate(30, allocator.kDefault);
+
+    {
+      auto best_fit_allocation =
+          dynamic_cast<BestFitAllocation*>(allocation2.get());
+      ASSERT_EQ(best_fit_allocation->ChunkIterator()->offset_, 80);
+    }
+    allocation2.reset();
+    allocation2 = allocator.Allocate(60, allocator.kDefault);
+
+    {
+      auto best_fit_allocation =
+          dynamic_cast<BestFitAllocation*>(allocation2.get());
+      ASSERT_EQ(best_fit_allocation->ChunkIterator()->offset_, 80);
+    }
+
+    allocation.reset();
+    allocation2.reset();
+
+    allocation = allocator.Allocate(80 + 60, allocator.kDefault);
+    {
+      auto best_fit_allocation =
+          dynamic_cast<BestFitAllocation*>(allocation.get());
+      ASSERT_EQ(best_fit_allocation->ChunkIterator()->offset_, 0);
+    }
+
+    allocation.reset();
+
+    allocation = allocator.Allocate(80, allocator.kDefault);
+    allocation2 = allocator.Allocate(60, allocator.kDefault);
+    allocation = nullptr;
+    allocation2 = nullptr;
+    allocation3 = nullptr;
+
+    ASSERT_EQ(allocator.NumFreeChunks(), 1U);
+  }
+}
+
+TEST(BestFitAllocator, test_concurrent_cpu_allocation) {
+  CPUAllocator allocator;
+  auto global_allocation =
+      allocator.Allocate(256UL * 1024 * 1024, allocator.kDefault);
+
+  std::unique_ptr<Allocator> best_fit_allocator(
+      new BestFitAllocator(global_allocation.get()));
+
+  LockedAllocator locked_allocator(std::move(best_fit_allocator));
+
+  auto th_main = [&] {
+    std::random_device dev;
+    std::default_random_engine engine(dev());
+    std::uniform_int_distribution<size_t> dist(1U, 1024U);
+
+    for (size_t i = 0; i < 128; ++i) {
+      size_t allocate_size = dist(engine);
+
+      auto allocation = locked_allocator.Allocate(
+          sizeof(size_t) * allocate_size, locked_allocator.kDefault);
+
+      size_t* data = reinterpret_cast<size_t*>(allocation->ptr());
+
+      for (size_t j = 0; j < allocate_size; ++j) {
+        data[j] = j;
+      }
+      std::this_thread::yield();
+
+      for (size_t j = 0; j < allocate_size; ++j) {
+        ASSERT_EQ(data[j], j);
+      }
+    }
+  };
+  {
+    std::vector<std::thread> threads;
+    for (size_t i = 0; i < 1024; ++i) {
+      threads.emplace_back(th_main);
+    }
+    for (auto& th : threads) {
+      th.join();
+    }
+  }
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/best_fit_allocator_test.cu

+// 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 <thread>  // NOLINT
+#include <vector>
+#include "gtest/gtest.h"
+#include "paddle/fluid/memory/allocation/best_fit_allocator.h"
+#include "paddle/fluid/memory/allocation/cuda_allocator.h"
+#include "paddle/fluid/memory/allocation/locked_allocator.h"
+#include "paddle/fluid/memory/memcpy.h"
+#include "paddle/fluid/platform/for_range.h"
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+struct ForEachFill {
+  size_t* ptr_;
+
+  explicit ForEachFill(size_t* ptr) : ptr_(ptr) {}
+
+  __device__ void operator()(size_t i) { ptr_[i] = i; }
+};
+
+TEST(BestFitAllocator, concurrent_cuda) {
+  CUDAAllocator allocator(platform::CUDAPlace(0));
+  // 256 MB
+  auto cuda_allocation =
+      allocator.Allocate(256U * 1024 * 1024, allocator.kDefault);
+  LockedAllocator concurrent_allocator(
+      std::unique_ptr<Allocator>(new BestFitAllocator(cuda_allocation.get())));
+
+  auto th_main = [&] {
+    std::random_device dev;
+    std::default_random_engine engine(dev());
+    std::uniform_int_distribution<size_t> dist(1U, 1024U);
+    platform::CUDAPlace gpu(0);
+    platform::CUDADeviceContext dev_ctx(gpu);
+    std::array<size_t, 1024> buf;
+    for (size_t i = 0; i < 128; ++i) {
+      size_t allocate_size = dist(engine);
+
+      auto allocation = concurrent_allocator.Allocate(
+          sizeof(size_t) * allocate_size, concurrent_allocator.kDefault);
+
+      size_t* data = reinterpret_cast<size_t*>(allocation->ptr());
+
+      ForEachFill fill(data);
+      platform::ForRange<platform::CUDADeviceContext> for_range(dev_ctx,
+                                                                allocate_size);
+      for_range(fill);
+
+      memory::Copy(platform::CPUPlace(), buf.data(), gpu, data,
+                   sizeof(size_t) * allocate_size, dev_ctx.stream());
+
+      dev_ctx.Wait();
+      for (size_t j = 0; j < allocate_size; ++j) {
+        ASSERT_EQ(buf[j], j);
+      }
+      allocation = nullptr;
+    }
+  };
+
+  {
+    std::vector<std::thread> threads;
+    for (size_t i = 0; i < 1024; ++i) {
+      threads.emplace_back(th_main);
+    }
+    for (auto& th : threads) {
+      th.join();
+    }
+  }
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/buffered_allocator.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/memory/allocation/buffered_allocator.h"
+#include <algorithm>
+#include <limits>
+#include <utility>
+#include "paddle/fluid/memory/allocation/allocation_with_underlying.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+BufferedAllocator::BufferedAllocator(std::unique_ptr<Allocator> &&allocator)
+    : underlying_allocator_(std::move(allocator)) {
+  PADDLE_ENFORCE_NOT_NULL(
+      underlying_allocator_,
+      "Underlying allocator of BufferedAllocator must be unmanaged");
+  if (underlying_allocator_->IsAllocThreadSafe()) {
+    mtx_.reset(new std::mutex());
+  }
+}
+
+BufferedAllocator::~BufferedAllocator() { FreeCache(-1UL); }
+
+void BufferedAllocator::FreeCache(size_t size) {
+  platform::LockGuardPtr<std::mutex> guard(mtx_);
+  if (UNLIKELY(size == 0)) return;
+  size_t cur = 0;
+  while (!allocations_.empty()) {  // free the largest
+    auto it = --allocations_.end();
+    cur += it->second->size();
+    delete it->second.release();
+    allocations_.erase(it);
+    if (cur >= size) return;
+  }
+}
+
+bool BufferedAllocator::IsAllocThreadSafe() const {
+  return this->underlying_allocator_->IsAllocThreadSafe();
+}
+void BufferedAllocator::Free(Allocation *allocation) {
+  platform::LockGuardPtr<std::mutex> guard(mtx_);
+  allocations_.emplace(allocation->size(), AllocationPtr(allocation));
+}
+Allocation *BufferedAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
+  {
+    platform::LockGuardPtr<std::mutex> guard(mtx_);
+    auto it = allocations_.lower_bound(size);
+    if (it != allocations_.end() && it->first < size * 2) {
+      AllocationPtr result(std::move(it->second));
+      allocations_.erase(it);
+      return new AllocationWithUnderlying(std::move(result));
+    }
+  }
+
+  try {
+    return new AllocationWithUnderlying(
+        underlying_allocator_->Allocate(size, attr));
+  } catch (BadAlloc &) {
+    FreeCache(size);
+    return new AllocationWithUnderlying(
+        underlying_allocator_->Allocate(size, attr));
+  }
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/buffered_allocator.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 <cstdint>
+#include <map>
+#include <memory>
+#include <vector>
+#include "paddle/fluid/memory/allocation/allocator.h"
+#include "paddle/fluid/platform/lock_guard_ptr.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// NOTE(zjl): BufferedAllocator maintains a memory pool to accelerate
+// memory allocation and reuse memory.
+// BufferedAllocator provides the same thread-safety level as
+// underlying_allocator_
+class BufferedAllocator : public Allocator {
+ public:
+  explicit BufferedAllocator(std::unique_ptr<Allocator> &&allocator);
+
+  ~BufferedAllocator();
+
+  bool IsAllocThreadSafe() const override;
+
+  // only used in unittest
+  inline void ClearCache() { FreeCache(-1UL); }
+
+ private:
+  void FreeCache(size_t size);
+
+ protected:
+  void Free(Allocation *allocation) override;
+  Allocation *AllocateImpl(size_t size, Allocator::Attr attr) override;
+
+ private:
+  std::unique_ptr<Allocator> underlying_allocator_;
+  std::multimap<size_t, AllocationPtr> allocations_;
+  std::unique_ptr<std::mutex> mtx_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/buffered_allocator_test.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/memory/allocation/buffered_allocator.h"
+#include <gtest/gtest.h>
+#include "paddle/fluid/memory/allocation/best_fit_allocator.h"
+#include "paddle/fluid/memory/allocation/cpu_allocator.h"
+#include "paddle/fluid/memory/allocation/locked_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+inline std::unique_ptr<BufferedAllocator> GetBufferedAllocator(
+    Allocation *allocation, bool thread_safe) {
+  std::unique_ptr<Allocator> allocator(new BestFitAllocator(allocation));
+  if (thread_safe) {
+    allocator.reset(new LockedAllocator(std::move(allocator)));
+  }
+
+  return std::unique_ptr<BufferedAllocator>(
+      new BufferedAllocator(std::move(allocator)));
+}
+
+TEST(buffered_allocator, thread_safety) {
+  std::unique_ptr<CPUAllocator> allocator(new CPUAllocator());
+  auto chunk = allocator->Allocate(1 << 20, allocator->kDefault);
+  {
+    auto buf_allocator = GetBufferedAllocator(chunk.get(), true);
+    ASSERT_EQ(buf_allocator->IsAllocThreadSafe(), true);
+  }
+
+  {
+    auto buf_allocator = GetBufferedAllocator(chunk.get(), false);
+    ASSERT_EQ(buf_allocator->IsAllocThreadSafe(), false);
+  }
+}
+
+class StubAllocation : public Allocation {
+ public:
+  using Allocation::Allocation;
+};
+
+class StubAllocator : public Allocator {
+ public:
+  void ResetCounter() {
+    construct_count_ = 0;
+    destruct_count_ = 0;
+  }
+
+  size_t GetAllocCount() const { return construct_count_; }
+
+  size_t GetFreeCount() const { return destruct_count_; }
+
+ protected:
+  void Free(Allocation *allocation) override {
+    auto *alloc = dynamic_cast<StubAllocation *>(allocation);
+    PADDLE_ENFORCE_NOT_NULL(alloc);
+    if (alloc->ptr()) delete[] static_cast<uint8_t *>(alloc->ptr());
+    ++destruct_count_;
+    delete allocation;
+  }
+  Allocation *AllocateImpl(size_t size, Allocator::Attr attr) override {
+    ++construct_count_;
+    if (size == 0) {
+      return new StubAllocation(nullptr, 0, platform::CPUPlace());
+    } else {
+      return new StubAllocation(new uint8_t[size], size, platform::CPUPlace());
+    }
+  }
+
+ private:
+  size_t construct_count_ = 0;
+  size_t destruct_count_ = 0;
+};
+
+constexpr size_t kZero = 0;
+constexpr size_t kOne = 1;
+constexpr size_t kTwo = 2;
+
+TEST(buffered_allocator, lazy_free) {
+  std::unique_ptr<StubAllocator> stub_allocator(new StubAllocator());
+  auto *underlying_allocator = stub_allocator.get();
+  std::unique_ptr<BufferedAllocator> allocator(
+      new BufferedAllocator(std::move(stub_allocator)));
+
+  {
+    underlying_allocator->ResetCounter();
+    auto x = allocator->Allocate(1025, allocator->kDefault);
+    ASSERT_EQ(underlying_allocator->GetAllocCount(), kOne);
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+    x = nullptr;
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+  }
+
+  {
+    underlying_allocator->ResetCounter();
+    auto x = allocator->Allocate(900, allocator->kDefault);
+    ASSERT_EQ(underlying_allocator->GetAllocCount(), kZero);
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+    auto y = allocator->Allocate(2048, allocator->kDefault);
+    ASSERT_EQ(underlying_allocator->GetAllocCount(), kOne);
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+    x = nullptr;
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+    y = nullptr;
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kZero);
+  }
+
+  {
+    underlying_allocator->ResetCounter();
+    allocator->ClearCache();
+    ASSERT_EQ(underlying_allocator->GetAllocCount(), kZero);
+    ASSERT_EQ(underlying_allocator->GetFreeCount(), kTwo);
+  }
+}
+
+TEST(buffered_allocator, garbage_collection) {
+  std::unique_ptr<CPUAllocator> cpu_allocator(new CPUAllocator());
+  auto chunk = cpu_allocator->Allocate(2048, cpu_allocator->kDefault);
+  auto allocator = GetBufferedAllocator(chunk.get(), false);
+  auto x1 = allocator->Allocate(1600, allocator->kDefault);
+  auto x2 = allocator->Allocate(400, allocator->kDefault);
+  x1 = nullptr;
+  x2 = nullptr;
+  auto x3 = allocator->Allocate(1600, allocator->kDefault);
+  ASSERT_NE(x3, nullptr);
+  ASSERT_NE(x3->ptr(), nullptr);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/conditional_allocator.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/memory/allocation/conditional_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+ConditionalAllocator& ConditionalAllocator::AddAllocator(
+    std::function<bool(size_t, Allocator::Attr)> func,
+    std::shared_ptr<Allocator> allocator) {
+  underlying_allocators_.emplace_back(std::move(func), std::move(allocator));
+  return *this;
+}
+
+bool ConditionalAllocator::IsAllocThreadSafe() const {
+  return std::all_of(underlying_allocators_.begin(),
+                     underlying_allocators_.end(),
+                     [](const AllocatorWithCond& allocatorWithCond) {
+                       return allocatorWithCond.second->IsAllocThreadSafe();
+                     });
+}
+
+Allocation* ConditionalAllocator::AllocateImpl(size_t size,
+                                               Allocator::Attr attr) {
+  for (auto& pair : underlying_allocators_) {
+    if (pair.first(size, attr)) {
+      return pair.second->Allocate(size, attr).release();
+    }
+  }
+  throw BadAlloc("No suitable allocator");
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/conditional_allocator.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 <utility>
+#include <vector>
+#include "paddle/fluid/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// A composite allocator who will dispatch the allocation request by registered
+// condition.
+//
+// For example:
+//
+// auto* cond_allocator = new ConditionalAllocator();
+// cond_allocator->AddAllocator([](size_t size, Attr attr){
+//   // if size > 10
+//   return size > 10;
+// }, allocator_a).AddAllocator([](size_t size, Attr attr){
+//   // elif attr is kDefault
+//   return attr == kDefault;
+// }, allocator_b).AddAllocator([](size_t size, Attr attr){
+//   // else
+//   return true;
+// }, allocator_c);
+class ConditionalAllocator : public Allocator {
+ public:
+  ConditionalAllocator() = default;
+
+  ConditionalAllocator& AddAllocator(std::function<bool(size_t, Attr)> func,
+                                     std::shared_ptr<Allocator> allocator);
+
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
+
+ private:
+  using AllocatorWithCond =
+      std::pair<std::function<bool(size_t, Attr)>, std::shared_ptr<Allocator>>;
+  std::vector<AllocatorWithCond> underlying_allocators_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/cpu_allocator.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/memory/allocation/cpu_allocator.h"
+#include <stdlib.h>
+#include <string>
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+CPUAllocation::CPUAllocation(void *ptr, size_t size)
+    : Allocation(ptr, size, platform::CPUPlace()) {}
+
+bool CPUAllocator::IsAllocThreadSafe() const { return true; }
+
+void CPUAllocator::Free(Allocation *allocation) {
+  PADDLE_ENFORCE_NOT_NULL(dynamic_cast<CPUAllocation *>(allocation));
+  free(allocation->ptr());
+  delete allocation;
+}
+
+Allocation *CPUAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
+  void *ptr;
+  auto status = posix_memalign(&ptr, kAlignment, size);
+  if (UNLIKELY(status) != 0) {
+    throw BadAlloc(string::Sprintf("Cannot allocate cpu memory %d. Errno is %d",
+                                   size, status));
+  }
+  return new CPUAllocation(ptr, size);
+}
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/cpu_allocator.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/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+// CPU system allocator and allocation.
+//
+// NOTE(yy): Should we just use `malloc` here since there is an
+// aligned_allocator.
+//
+// NOTE(yy): It is no need to use `BestFitAllocator` in CPU. We can import
+// an open-sourced allocator into Paddle.
+class CPUAllocator;
+class CPUAllocation : public Allocation {
+ public:
+  CPUAllocation(void* ptr, size_t size);
+};
+
+class CPUAllocator : public Allocator {
+ public:
+  constexpr static size_t kAlignment = 64u;
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void Free(Allocation* allocation) override;
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
+};
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/cuda_allocator.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/memory/allocation/cuda_allocator.h"
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <string>
+#include "paddle/fluid/platform/cuda_device_guard.h"
+#include "paddle/fluid/platform/gpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+bool CUDAAllocator::IsAllocThreadSafe() const { return true; }
+void CUDAAllocator::Free(Allocation* allocation) {
+  platform::CUDADeviceGuard guard(place_.device);
+  auto* cuda_allocation = dynamic_cast<CUDAAllocation*>(allocation);
+  PADDLE_ENFORCE_NOT_NULL(cuda_allocation);
+  PADDLE_ENFORCE_EQ(boost::get<platform::CUDAPlace>(cuda_allocation->place()),
+                    place_);
+  PADDLE_ENFORCE(cudaFree(allocation->ptr()));
+  delete allocation;
+}
+Allocation* CUDAAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
+  platform::CUDADeviceGuard guard(place_.device);
+  void* ptr;
+  auto status = cudaMalloc(&ptr, size);
+  if (UNLIKELY(status != cudaSuccess)) {
+    throw BadAlloc(string::Sprintf(
+        "Cannot allocate %d on GPU %d, cuda status %d, %s", size, place_.device,
+        status, cudaGetErrorString(status)));
+  }
+  return new CUDAAllocation(ptr, size, platform::Place(place_));
+}
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/cuda_allocator.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/memory/allocation/allocator.h"
+#include "paddle/fluid/platform/place.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// CUDA System allocator and allocation.
+// Just a flag type.
+class CUDAAllocation : public Allocation {
+ public:
+  using Allocation::Allocation;
+};
+
+class CUDAAllocator : public Allocator {
+ public:
+  explicit CUDAAllocator(const platform::CUDAPlace& place) : place_(place) {}
+  explicit CUDAAllocator(const platform::Place& place)
+      : place_(boost::get<platform::CUDAPlace>(place)) {}
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void Free(Allocation* allocation) override;
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
+
+ private:
+  platform::CUDAPlace place_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/legacy_allocator.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/memory/allocation/legacy_allocator.h"
+#include <string>
+#include "glog/logging.h"
+#include "paddle/fluid/memory/detail/buddy_allocator.h"
+#include "paddle/fluid/memory/detail/system_allocator.h"
+#include "paddle/fluid/platform/gpu_info.h"
+#include "paddle/fluid/string/printf.h"
+
+DEFINE_bool(init_allocated_mem, false,
+            "It is a mistake that the values of the memory allocated by "
+            "BuddyAllocator are always zeroed in some op's implementation. "
+            "To find this error in time, we use init_allocated_mem to indicate "
+            "that initializing the allocated memory with a small value "
+            "during unit testing.");
+DECLARE_double(fraction_of_gpu_memory_to_use);
+
+namespace paddle {
+namespace memory {
+namespace legacy {
+template <typename Place>
+void *Alloc(const Place &place, size_t size);
+
+template <typename Place>
+void Free(const Place &place, void *p);
+
+template <typename Place>
+size_t Used(const Place &place);
+
+struct Usage : public boost::static_visitor<size_t> {
+  size_t operator()(const platform::CPUPlace &cpu) const;
+  size_t operator()(const platform::CUDAPlace &gpu) const;
+  size_t operator()(const platform::CUDAPinnedPlace &cuda_pinned) const;
+};
+
+size_t memory_usage(const platform::Place &p);
+
+using BuddyAllocator = detail::BuddyAllocator;
+
+BuddyAllocator *GetCPUBuddyAllocator() {
+  // We tried thread_local for inference::RNN1 model, but that not works much
+  // for multi-thread test.
+  static std::once_flag init_flag;
+  static detail::BuddyAllocator *a = nullptr;
+
+  std::call_once(init_flag, []() {
+    a = new detail::BuddyAllocator(
+        std::unique_ptr<detail::SystemAllocator>(new detail::CPUAllocator),
+        platform::CpuMinChunkSize(), platform::CpuMaxChunkSize());
+  });
+
+  return a;
+}
+
+// We compared the NaiveAllocator with BuddyAllocator in CPU memory allocation,
+// seems they are almost the same overhead.
+struct NaiveAllocator {
+  void *Alloc(size_t size) { return malloc(size); }
+
+  void Free(void *p) {
+    PADDLE_ENFORCE(p);
+    free(p);
+  }
+
+  static NaiveAllocator *Instance() {
+    static NaiveAllocator x;
+    return &x;
+  }
+
+ private:
+  std::mutex lock_;
+};
+
+template <>
+void *Alloc<platform::CPUPlace>(const platform::CPUPlace &place, size_t size) {
+  VLOG(10) << "Allocate " << size << " bytes on " << platform::Place(place);
+  void *p = GetCPUBuddyAllocator()->Alloc(size);
+  if (FLAGS_init_allocated_mem) {
+    memset(p, 0xEF, size);
+  }
+  VLOG(100) << "  pointer=" << p;
+  return p;
+}
+
+template <>
+void Free<platform::CPUPlace>(const platform::CPUPlace &place, void *p) {
+  VLOG(10) << "Free pointer=" << p << " on " << platform::Place(place);
+  GetCPUBuddyAllocator()->Free(p);
+}
+
+template <>
+size_t Used<platform::CPUPlace>(const platform::CPUPlace &place) {
+  return GetCPUBuddyAllocator()->Used();
+}
+
+#ifdef PADDLE_WITH_CUDA
+BuddyAllocator *GetGPUBuddyAllocator(int gpu_id) {
+  static std::once_flag init_flag;
+  static detail::BuddyAllocator **a_arr = nullptr;
+
+  std::call_once(init_flag, [gpu_id]() {
+    int gpu_num = platform::GetCUDADeviceCount();
+    PADDLE_ENFORCE(gpu_id < gpu_num, "gpu_id:%d should < gpu_num:%d", gpu_id,
+                   gpu_num);
+
+    a_arr = new BuddyAllocator *[gpu_num];
+    for (int i = 0; i < gpu_num; i++) {
+      a_arr[i] = nullptr;
+      platform::SetDeviceId(i);
+      a_arr[i] = new BuddyAllocator(
+          std::unique_ptr<detail::SystemAllocator>(new detail::GPUAllocator(i)),
+          platform::GpuMinChunkSize(), platform::GpuMaxChunkSize());
+
+      VLOG(100) << "\n\nNOTE: each GPU device use "
+                << FLAGS_fraction_of_gpu_memory_to_use * 100
+                << "% of GPU memory.\n"
+                << "You can set GFlags environment variable '"
+                << "FLAGS_fraction_of_gpu_memory_to_use"
+                << "' to change the fraction of GPU usage.\n\n";
+    }
+  });
+
+  platform::SetDeviceId(gpu_id);
+  return a_arr[gpu_id];
+}
+#endif
+
+template <>
+size_t Used<platform::CUDAPlace>(const platform::CUDAPlace &place) {
+#ifdef PADDLE_WITH_CUDA
+  return GetGPUBuddyAllocator(place.device)->Used();
+#else
+  PADDLE_THROW("'CUDAPlace' is not supported in CPU only device.");
+#endif
+}
+
+template <>
+void *Alloc<platform::CUDAPlace>(const platform::CUDAPlace &place,
+                                 size_t size) {
+#ifdef PADDLE_WITH_CUDA
+  auto *buddy_allocator = GetGPUBuddyAllocator(place.device);
+  auto *ptr = buddy_allocator->Alloc(size);
+  if (ptr == nullptr) {
+    int cur_dev = platform::GetCurrentDeviceId();
+    platform::SetDeviceId(place.device);
+    size_t avail, total;
+    platform::GpuMemoryUsage(&avail, &total);
+    LOG(WARNING) << "Cannot allocate " << string::HumanReadableSize(size)
+                 << " in GPU " << place.device << ", available "
+                 << string::HumanReadableSize(avail);
+    LOG(WARNING) << "total " << total;
+    LOG(WARNING) << "GpuMinChunkSize "
+                 << string::HumanReadableSize(
+                        buddy_allocator->GetMinChunkSize());
+    LOG(WARNING) << "GpuMaxChunkSize "
+                 << string::HumanReadableSize(
+                        buddy_allocator->GetMaxChunkSize());
+    LOG(WARNING) << "GPU memory used: "
+                 << string::HumanReadableSize(Used<platform::CUDAPlace>(place));
+    platform::SetDeviceId(cur_dev);
+  }
+  if (FLAGS_init_allocated_mem) {
+    cudaMemset(ptr, 0xEF, size);
+  }
+  return ptr;
+#else
+  PADDLE_THROW("'CUDAPlace' is not supported in CPU only device.");
+#endif
+}
+
+template <>
+void Free<platform::CUDAPlace>(const platform::CUDAPlace &place, void *p) {
+#ifdef PADDLE_WITH_CUDA
+  GetGPUBuddyAllocator(place.device)->Free(p);
+#else
+  PADDLE_THROW("'CUDAPlace' is not supported in CPU only device.");
+#endif
+}
+
+#ifdef PADDLE_WITH_CUDA
+BuddyAllocator *GetCUDAPinnedBuddyAllocator() {
+  static std::once_flag init_flag;
+  static BuddyAllocator *ba = nullptr;
+
+  std::call_once(init_flag, []() {
+    ba = new BuddyAllocator(std::unique_ptr<detail::SystemAllocator>(
+                                new detail::CUDAPinnedAllocator),
+                            platform::CUDAPinnedMinChunkSize(),
+                            platform::CUDAPinnedMaxChunkSize());
+  });
+
+  return ba;
+}
+#endif
+
+template <>
+size_t Used<platform::CUDAPinnedPlace>(const platform::CUDAPinnedPlace &place) {
+#ifdef PADDLE_WITH_CUDA
+  return GetCUDAPinnedBuddyAllocator()->Used();
+#else
+  PADDLE_THROW("'CUDAPinnedPlace' is not supported in CPU only device.");
+#endif
+}
+
+template <>
+void *Alloc<platform::CUDAPinnedPlace>(const platform::CUDAPinnedPlace &place,
+                                       size_t size) {
+#ifdef PADDLE_WITH_CUDA
+  auto *buddy_allocator = GetCUDAPinnedBuddyAllocator();
+  void *ptr = buddy_allocator->Alloc(size);
+
+  if (ptr == nullptr) {
+    LOG(WARNING) << "cudaMallocHost Cannot allocate " << size
+                 << " bytes in CUDAPinnedPlace";
+  }
+  if (FLAGS_init_allocated_mem) {
+    memset(ptr, 0xEF, size);
+  }
+  return ptr;
+#else
+  PADDLE_THROW("'CUDAPinnedPlace' is not supported in CPU only device.");
+#endif
+}
+
+template <>
+void Free<platform::CUDAPinnedPlace>(const platform::CUDAPinnedPlace &place,
+                                     void *p) {
+#ifdef PADDLE_WITH_CUDA
+  GetCUDAPinnedBuddyAllocator()->Free(p);
+#else
+  PADDLE_THROW("'CUDAPinnedPlace' is not supported in CPU only device.");
+#endif
+}
+
+struct AllocVisitor : public boost::static_visitor<void *> {
+  inline explicit AllocVisitor(size_t size) : size_(size) {}
+
+  template <typename Place>
+  inline void *operator()(const Place &place) const {
+    return Alloc<Place>(place, size_);
+  }
+
+ private:
+  size_t size_;
+};
+
+struct FreeVisitor : public boost::static_visitor<void> {
+  inline explicit FreeVisitor(void *ptr) : ptr_(ptr) {}
+
+  template <typename Place>
+  inline void operator()(const Place &place) const {
+    Free<Place>(place, ptr_);
+  }
+
+ private:
+  void *ptr_;
+};
+
+size_t Usage::operator()(const platform::CPUPlace &cpu) const {
+  return Used(cpu);
+}
+
+size_t Usage::operator()(const platform::CUDAPlace &gpu) const {
+#ifdef PADDLE_WITH_CUDA
+  return Used(gpu);
+#else
+  PADDLE_THROW("'CUDAPlace' is not supported in CPU only device.");
+#endif
+}
+
+size_t Usage::operator()(const platform::CUDAPinnedPlace &cuda_pinned) const {
+#ifdef PADDLE_WITH_CUDA
+  return Used(cuda_pinned);
+#else
+  PADDLE_THROW("'CUDAPinnedPlace' is not supported in CPU only device.");
+#endif
+}
+}  // namespace legacy
+
+namespace allocation {
+
+Allocation *LegacyAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
+  void *ptr = boost::apply_visitor(legacy::AllocVisitor(size), place_);
+  return new Allocation(ptr, size, place_);
+}
+
+void LegacyAllocator::Free(Allocation *allocation) {
+  boost::apply_visitor(legacy::FreeVisitor(allocation->ptr()),
+                       allocation->place());
+  delete allocation;
+}
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/legacy_allocator.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/memory/allocation/allocator.h"
+#include "paddle/fluid/platform/place.h"
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class LegacyAllocatorPrivate;
+class LegacyAllocator : public Allocator {
+ public:
+  explicit LegacyAllocator(const platform::Place &p) : place_(p) {}
+
+ protected:
+  Allocation *AllocateImpl(size_t size, Allocator::Attr attr) override;
+  void Free(Allocation *allocation) override;
+
+ private:
+  platform::Place place_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/locked_allocator.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 <mutex>   // NOLINT
+#include <thread>  // NOLINT
+#include "paddle/fluid/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// A allocator to make underlying allocator thread safe.
+class LockedAllocator : public Allocator {
+ public:
+  explicit LockedAllocator(std::unique_ptr<Allocator> &&underlying_allocator);
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void Free(Allocation *allocation) override;
+  Allocation *AllocateImpl(size_t size, Allocator::Attr attr) override;
+
+ private:
+  std::unique_ptr<Allocator> underlying_allocator_;
+  std::unique_ptr<std::mutex> mtx_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle