Merge branch 'windows/build' into windows/online

test=develop

.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

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_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/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/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.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/locked_allocator.h"
+#include <mutex>  // NOLINT
+#include "paddle/fluid/memory/allocation/allocation_with_underlying.h"
+#include "paddle/fluid/platform/lock_guard_ptr.h"
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+bool LockedAllocator::IsAllocThreadSafe() const { return true; }
+
+LockedAllocator::LockedAllocator(
+    std::unique_ptr<Allocator> &&underlying_allocator)
+    : underlying_allocator_(std::move(underlying_allocator)) {
+  PADDLE_ENFORCE_NOT_NULL(underlying_allocator_);
+  if (!underlying_allocator_->IsAllocThreadSafe()) {
+    mtx_.reset(new std::mutex());
+  }
+}
+void LockedAllocator::Free(Allocation *allocation) {
+  {
+    platform::LockGuardPtr<std::mutex> guard(mtx_);
+    reinterpret_cast<AllocationWithUnderlying *>(allocation)
+        ->allocation_.reset();  // Destroy inner allocation
+  }
+  delete allocation;
+}
+Allocation *LockedAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
+  platform::LockGuardPtr<std::mutex> guard(mtx_);
+  return new AllocationWithUnderlying(
+      underlying_allocator_->Allocate(size, attr));
+}
+}  // 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

paddle/fluid/memory/allocation/pinned_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/pinned_allocator.h"
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+bool CPUPinnedAllocator::IsAllocThreadSafe() const { return true; }
+void CPUPinnedAllocator::Free(Allocation *allocation) {
+  PADDLE_ENFORCE_NOT_NULL(dynamic_cast<CPUPinnedAllocation *>(allocation));
+  PADDLE_ENFORCE(cudaFreeHost(allocation->ptr()));
+  delete allocation;
+}
+Allocation *CPUPinnedAllocator::AllocateImpl(size_t size,
+                                             Allocator::Attr attr) {
+  // PADDLE_ENFORCE_EQ(
+  //    attr, kCrossDevice,
+  //    "CPUPinnedAllocator should be used for Cross-Device Communication");
+
+  void *ptr;
+  PADDLE_ENFORCE(cudaMallocHost(&ptr, size));
+  return new CPUPinnedAllocation(ptr, size);
+}
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/pinned_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 {
+
+// Allocator uses `cudaMallocHost`
+class CPUPinnedAllocation : public Allocation {
+ public:
+  CPUPinnedAllocation(void *ptr, size_t size)
+      : Allocation(ptr, size, platform::CUDAPinnedPlace()) {}
+};
+
+class CPUPinnedAllocator : public Allocator {
+ public:
+  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/retry_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/retry_allocator.h"
+#include "paddle/fluid/memory/allocation/allocation_with_underlying.h"
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+bool RetryAllocator::IsAllocThreadSafe() const {
+  return underlying_allocator_->IsAllocThreadSafe();
+}
+
+void RetryAllocator::Free(Allocation* allocation) {
+  // Delete underlying allocation first.
+  reinterpret_cast<AllocationWithUnderlying*>(allocation)->allocation_.reset();
+  {
+    // notify all waited allocators, they can try to allocate memory after free.
+    std::lock_guard<std::mutex> lock(mutex_);
+    cv_.notify_all();
+  }
+  delete allocation;
+}
+
+Allocation* RetryAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
+  auto alloc_func = [&, this]() {
+    return new AllocationWithUnderlying(
+        underlying_allocator_->Allocate(size, attr));
+  };
+  // In fact, we can unify the code of allocation success and failure
+  // But it would add lock even when allocation success at the first time
+  try {
+    return alloc_func();
+  } catch (BadAlloc& bad_alloc) {
+    {
+      // We can just write allocation retry inside the predicate function of
+      // wait_until
+      // But it needs to acquire the lock when executing predicate function
+      // For better performance, we use loop here
+      auto end_time = std::chrono::high_resolution_clock::now() + retry_time_;
+      auto wait_until = [&, this] {
+        std::unique_lock<std::mutex> lock(mutex_);
+        return cv_.wait_until(lock, end_time);
+      };
+      while (wait_until() != std::cv_status::timeout) {
+        try {
+          return alloc_func();
+        } catch (BadAlloc& ex) {
+          bad_alloc = ex;
+        } catch (...) {
+          throw;
+        }
+      }
+
+      throw;  // rethrow the original exception or throw the internal bad_alloc
+    }
+  } catch (...) {
+    throw;
+  }
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/retry_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 <chrono>              // NOLINT
+#include <condition_variable>  // NOLINT
+#include <memory>
+#include <mutex>  // NOLINT
+#include "paddle/fluid/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class RetryAllocator;
+
+class RetryAllocator : public Allocator {
+ public:
+  RetryAllocator(std::unique_ptr<Allocator>&& allocator, size_t retry_ms)
+      : underlying_allocator_(std::move(allocator)), retry_time_(retry_ms) {
+    EnforceCheck();
+  }
+
+  bool IsAllocThreadSafe() const override;
+
+ private:
+  void EnforceCheck() {
+    PADDLE_ENFORCE_NOT_NULL(
+        underlying_allocator_.get(),
+        "UnderlyingAllocator of RetryAllocator must be UnmanagedAllocator");
+    PADDLE_ENFORCE(underlying_allocator_->IsAllocThreadSafe(),
+                   "UnderlyingAllocator of RetryAllocator must be thread-safe");
+  }
+
+ protected:
+  void Free(Allocation* allocation) override;
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
+
+ private:
+  std::unique_ptr<Allocator> underlying_allocator_;
+  std::chrono::milliseconds retry_time_;
+  std::mutex mutex_;
+  std::condition_variable cv_;
+
+  // For debug, We can add an atomic integer to record how many memory sizes are
+  // waited to allocate
+  // std::atomic<size_t> waited_allocate_size_{0};
+
+  friend class RetryAllocation;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/retry_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/retry_allocator.h"
+#include <algorithm>
+#include <chrono>              // NOLINT
+#include <condition_variable>  // NOLINT
+#include <mutex>               // NOLINT
+#include <thread>              // NOLINT
+#include <vector>
+#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 {
+
+TEST(RetryAllocator, RetryAllocator) {
+  CPUAllocator cpu_allocator;
+
+  size_t size = (1 << 20);
+  auto cpu_allocation = cpu_allocator.Allocate(size, cpu_allocator.kDefault);
+
+  std::unique_ptr<BestFitAllocator> best_fit_allocator(
+      new BestFitAllocator(cpu_allocation.get()));
+  std::unique_ptr<LockedAllocator> locked_allocator(
+      new LockedAllocator(std::move(best_fit_allocator)));
+
+  size_t thread_num = 32;
+  size_t sleep_time = 40;
+  size_t extra_time = 2;
+
+  // Reserve to perform more tests in the future
+  std::vector<std::shared_ptr<Allocator>> allocators;
+  {
+    std::unique_ptr<BestFitAllocator> best_fit_allocator(
+        new BestFitAllocator(cpu_allocation.get()));
+    std::unique_ptr<LockedAllocator> locked_allocator(
+        new LockedAllocator(std::move(best_fit_allocator)));
+    allocators.push_back(std::make_shared<RetryAllocator>(
+        std::move(locked_allocator),
+        (thread_num - 1) * (sleep_time + extra_time)));
+  }
+
+  for (auto &allocator : allocators) {
+    std::vector<std::thread> threads(thread_num);
+    std::vector<void *> addresses(threads.size(), nullptr);
+
+    std::mutex mutex;
+    std::condition_variable cv;
+    bool flag = false;
+
+    for (size_t i = 0; i < threads.size(); ++i) {
+      threads[i] = std::thread([&, i]() {
+        {
+          std::unique_lock<std::mutex> lock(mutex);
+          cv.wait(lock, [&] { return flag; });
+        }
+
+        auto ret = allocator->Allocate(size - 1);
+        addresses[i] = ret->ptr();
+        std::this_thread::sleep_for(std::chrono::milliseconds(sleep_time));
+      });
+    }
+
+    {
+      std::lock_guard<std::mutex> lock(mutex);
+      flag = true;
+      cv.notify_all();
+    }
+
+    for (auto &th : threads) {
+      th.join();
+    }
+
+    void *val = cpu_allocation->ptr();
+    bool is_all_equal = std::all_of(addresses.begin(), addresses.end(),
+                                    [val](void *p) { return p == val; });
+    ASSERT_TRUE(is_all_equal);
+  }
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/zero_size_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/zero_size_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+bool ZeroSizeAllocator::IsAllocThreadSafe() const {
+  return underlying_allocator_->IsAllocThreadSafe();
+}
+
+Allocation *ZeroSizeAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
+  if (size == 0) {
+    return new ZeroSizeAllocation(place_);
+  } else {
+    return underlying_allocator_->Allocate(size, attr).release();
+  }
+}
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/zero_size_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 <utility>
+#include "paddle/fluid/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// The allocator handles the request's size is zero. Allocator will always
+// return an allocation even the request size is zero. However, the
+// allocation.ptr() is nullptr
+class ZeroSizeAllocation : public Allocation {
+ public:
+  explicit ZeroSizeAllocation(const platform::Place& p)
+      : Allocation(nullptr, 0, p) {}
+};
+
+class ZeroSizeAllocator : public Allocator {
+ public:
+  ZeroSizeAllocator(std::shared_ptr<Allocator> underlying_allocator,
+                    const platform::Place& p)
+      : underlying_allocator_(std::move(underlying_allocator)), place_(p) {}
+
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
+
+ private:
+  std::shared_ptr<Allocator> underlying_allocator_;
+  const platform::Place& place_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/detail/system_allocator.cc

@@ -30,12 +30,7 @@ limitations under the License. */
 #include "paddle/fluid/platform/enforce.h"
 #include "paddle/fluid/platform/gpu_info.h"
 
-// If use_pinned_memory is true, CPUAllocator calls mlock, which
-// returns pinned and locked memory as staging areas for data exchange
-// between host and device.  Allocates too much would reduce the amount
-// of memory available to the system for paging.  So, by default, we
-// should set false to use_pinned_memory.
-DEFINE_bool(use_pinned_memory, true, "If set, allocate cpu pinned memory.");
+DECLARE_bool(use_pinned_memory);
 DECLARE_double(fraction_of_gpu_memory_to_use);
 namespace paddle {
 namespace memory {

paddle/fluid/memory/malloc.cc

@@ -12,221 +12,22 @@ 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/malloc.h"
 #include <string>
 #include <vector>
-
-#include "paddle/fluid/memory/malloc.h"
-
-#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);
-
+#include "paddle/fluid/memory/allocation/allocator_facade.h"
+#include "paddle/fluid/memory/allocation/allocator_strategy.h"
+#include "paddle/fluid/platform/place.h"
 namespace paddle {
 namespace memory {
-
-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>(platform::CPUPlace place, size_t size) {
-  VLOG(100) << "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>(platform::CPUPlace place, void* p) {
-  VLOG(100) << "Free pointer=" << p << " on " << platform::Place(place);
-  GetCPUBuddyAllocator()->Free(p);
-}
-
-template <>
-size_t Used<platform::CPUPlace>(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];
-}
-
-template <>
-size_t Used<platform::CUDAPlace>(platform::CUDAPlace place) {
-  return GetGPUBuddyAllocator(place.device)->Used();
-}
-
-template <>
-void* Alloc<platform::CUDAPlace>(platform::CUDAPlace place, size_t size) {
-  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;
-}
-
-template <>
-void Free<platform::CUDAPlace>(platform::CUDAPlace place, void* p) {
-  GetGPUBuddyAllocator(place.device)->Free(p);
-}
-
-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;
-}
-
-template <>
-size_t Used<platform::CUDAPinnedPlace>(platform::CUDAPinnedPlace place) {
-  return GetCUDAPinnedBuddyAllocator()->Used();
-}
-
-template <>
-void* Alloc<platform::CUDAPinnedPlace>(platform::CUDAPinnedPlace place,
-                                       size_t size) {
-  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;
-}
-
-template <>
-void Free<platform::CUDAPinnedPlace>(platform::CUDAPinnedPlace place, void* p) {
-  GetCUDAPinnedBuddyAllocator()->Free(p);
-}
-#endif
-
-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
+std::shared_ptr<Allocation> AllocShared(const platform::Place& place,
+                                        size_t size, Allocator::Attr attr) {
+  return allocation::AllocatorFacade::Instance().AllocShared(place, size, attr);
 }
 
-size_t memory_usage(const platform::Place& p) {
-  return boost::apply_visitor(Usage(), p);
+AllocationPtr Alloc(const platform::Place& place, size_t size,
+                    Allocator::Attr attr) {
+  return allocation::AllocatorFacade::Instance().Alloc(place, size, attr);
 }
 
 }  // namespace memory

paddle/fluid/memory/malloc.h

@@ -14,91 +14,21 @@ 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 {
+using allocation::Allocation;
+using allocation::Allocator;
+using allocation::AllocationPtr;
 
-/**
- * \brief   Allocate memory block in one place.
- *
- * \param[in]  place  Allocation place (CPU or GPU).
- * \param[in]  size   Allocation size.
- *
- * \return  Allocated memory block address.
- *
- * \note    If return nullptr, it indicates memory allocation failed
- *          because insufficient memory in current system. When Alloc
- *          function is invoked, you must check the returned memory
- *          address is valid or not.
- */
-template <typename Place>
-void* Alloc(Place place, size_t size);
-
-/**
- * \brief   Free memory block in one place.
- *
- * \param[in]  place  Allocation place (CPU or GPU).
- * \param[in]  ptr    Memory block address to free.
- *
- */
-template <typename Place>
-void Free(Place place, void* ptr);
-
-/**
- * \brief   Total size of used memory in one place.
- *
- * \param[in]  place  Allocation place (CPU or GPU).
- *
- */
-template <typename Place>
-size_t Used(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);
-
-/**
- * \brief   Free memory block in one place.
- *
- * \note    In some cases, custom deleter is used to
- *          deallocate the memory automatically for
- *          std::unique_ptr<T> in tensor.h.
- *
- */
-template <typename T, typename Place>
-class PODDeleter {
-  static_assert(std::is_pod<T>::value, "T must be POD");
-
- public:
-  explicit PODDeleter(Place place) : place_(place) {}
-  void operator()(T* ptr) { Free(place_, static_cast<void*>(ptr)); }
-
- private:
-  Place place_;
-};
-
-/**
- * \brief   Free memory block in one place does not meet POD
- *
- * \note    In some cases, custom deleter is used to
- *          deallocate the memory automatically for
- *          std::unique_ptr<T> in tensor.h.
- *
- */
-template <typename T, typename Place>
-class PlainDeleter {
- public:
-  explicit PlainDeleter(Place place) : place_(place) {}
-  void operator()(T* ptr) { Free(place_, reinterpret_cast<void*>(ptr)); }
+extern std::shared_ptr<Allocation> AllocShared(
+    const platform::Place& place, size_t size,
+    Allocator::Attr attr = Allocator::kDefault);
 
- private:
-  Place place_;
-};
+extern AllocationPtr Alloc(const platform::Place& place, size_t size,
+                           Allocator::Attr attr = Allocator::kDefault);
 
 }  // namespace memory
 }  // namespace paddle

paddle/fluid/memory/malloc_test.cc

-/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License. */
-
-#include "paddle/fluid/memory/malloc.h"
-
-#include <unordered_map>
-
-#include "gtest/gtest.h"
-#include "paddle/fluid/memory/detail/memory_block.h"
-#include "paddle/fluid/platform/cpu_info.h"
-#include "paddle/fluid/platform/gpu_info.h"
-#include "paddle/fluid/platform/place.h"
-
-inline bool is_aligned(void const *p) {
-  return 0 == (reinterpret_cast<uintptr_t>(p) & 0x3);
-}
-
-size_t align(size_t size, paddle::platform::CPUPlace place) {
-  size += sizeof(paddle::memory::detail::MemoryBlock::Desc);
-  size_t alignment = paddle::platform::CpuMinChunkSize();
-  size_t remaining = size % alignment;
-  return remaining == 0 ? size : size + (alignment - remaining);
-}
-
-TEST(BuddyAllocator, CPUAllocation) {
-  void *p = nullptr;
-
-  EXPECT_EQ(p, nullptr);
-
-  paddle::platform::CPUPlace cpu;
-  p = paddle::memory::Alloc(cpu, 4096);
-
-  EXPECT_NE(p, nullptr);
-
-  paddle::platform::Place place = cpu;
-  EXPECT_EQ(paddle::memory::Used(cpu), paddle::memory::memory_usage(place));
-
-  paddle::memory::Free(cpu, p);
-}
-
-TEST(BuddyAllocator, CPUMultAlloc) {
-  paddle::platform::CPUPlace cpu;
-
-  std::unordered_map<void *, size_t> ps;
-
-  size_t total_size = paddle::memory::Used(cpu);
-  EXPECT_EQ(total_size, 0UL);
-
-  for (auto size :
-       {0, 128, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304}) {
-    ps[paddle::memory::Alloc(cpu, size)] = size;
-
-    // Buddy Allocator doesn't manage too large memory chunk
-    if (paddle::memory::Used(cpu) == total_size) continue;
-
-    size_t aligned_size = align(size, cpu);
-    total_size += aligned_size;
-    EXPECT_EQ(total_size, paddle::memory::Used(cpu));
-  }
-
-  for (auto p : ps) {
-    EXPECT_EQ(is_aligned(p.first), true);
-    paddle::memory::Free(cpu, p.first);
-
-    // Buddy Allocator doesn't manage too large memory chunk
-    if (paddle::memory::Used(cpu) == total_size) continue;
-
-    size_t aligned_size = align(p.second, cpu);
-    total_size -= aligned_size;
-    EXPECT_EQ(total_size, paddle::memory::Used(cpu));
-  }
-}
-
-#ifdef PADDLE_WITH_CUDA
-
-size_t align(size_t size, paddle::platform::CUDAPlace place) {
-  size += sizeof(paddle::memory::detail::MemoryBlock::Desc);
-  size_t alignment = paddle::platform::GpuMinChunkSize();
-  size_t remaining = size % alignment;
-  return remaining == 0 ? size : size + (alignment - remaining);
-}
-
-TEST(BuddyAllocator, GPUAllocation) {
-  void *p = nullptr;
-
-  EXPECT_EQ(p, nullptr);
-
-  paddle::platform::CUDAPlace gpu(0);
-  p = paddle::memory::Alloc(gpu, 4096);
-
-  EXPECT_NE(p, nullptr);
-
-  paddle::platform::Place place = gpu;
-  EXPECT_EQ(paddle::memory::Used(gpu), paddle::memory::memory_usage(place));
-
-  paddle::memory::Free(gpu, p);
-}
-
-TEST(BuddyAllocator, GPUMultAlloc) {
-  paddle::platform::CUDAPlace gpu;
-
-  std::unordered_map<void *, size_t> ps;
-
-  size_t total_size = paddle::memory::Used(gpu);
-  EXPECT_EQ(total_size, 0UL);
-
-  for (auto size :
-       {0, 128, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304}) {
-    ps[paddle::memory::Alloc(gpu, size)] = size;
-
-    // Buddy Allocator doesn't manage too large memory chunk
-    if (paddle::memory::Used(gpu) == total_size) continue;
-
-    size_t aligned_size = align(size, gpu);
-    total_size += aligned_size;
-    EXPECT_EQ(total_size, paddle::memory::Used(gpu));
-  }
-
-  for (auto p : ps) {
-    EXPECT_EQ(is_aligned(p.first), true);
-    paddle::memory::Free(gpu, p.first);
-
-    // Buddy Allocator doesn't manage too large memory chunk
-    if (paddle::memory::Used(gpu) == total_size) continue;
-
-    size_t aligned_size = align(p.second, gpu);
-    total_size -= aligned_size;
-    EXPECT_EQ(total_size, paddle::memory::Used(gpu));
-  }
-}
-
-size_t align(size_t size, paddle::platform::CUDAPinnedPlace place) {
-  size += sizeof(paddle::memory::detail::MemoryBlock::Desc);
-  size_t alignment = paddle::platform::CUDAPinnedMinChunkSize();
-  size_t remaining = size % alignment;
-  return remaining == 0 ? size : size + (alignment - remaining);
-}
-
-TEST(BuddyAllocator, CUDAPinnedAllocator) {
-  void *p = nullptr;
-
-  EXPECT_EQ(p, nullptr);
-
-  paddle::platform::CUDAPinnedPlace cpu;
-  p = paddle::memory::Alloc(cpu, 4096);
-
-  EXPECT_NE(p, nullptr);
-
-  paddle::platform::Place place = cpu;
-  EXPECT_EQ(paddle::memory::Used(cpu), paddle::memory::memory_usage(place));
-
-  paddle::memory::Free(cpu, p);
-}
-
-TEST(BuddyAllocator, CUDAPinnedMultAllocator) {
-  paddle::platform::CUDAPinnedPlace cpu;
-
-  std::unordered_map<void *, size_t> ps;
-
-  size_t total_size = paddle::memory::Used(cpu);
-  EXPECT_EQ(total_size, 0UL);
-
-  for (auto size :
-       {0, 128, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304}) {
-    ps[paddle::memory::Alloc(cpu, size)] = size;
-
-    // Buddy Allocator doesn't manage too large memory chunk
-    if (paddle::memory::Used(cpu) == total_size) continue;
-
-    size_t aligned_size = align(size, cpu);
-    total_size += aligned_size;
-    EXPECT_EQ(total_size, paddle::memory::Used(cpu));
-  }
-
-  for (auto p : ps) {
-    EXPECT_EQ(is_aligned(p.first), true);
-    paddle::memory::Free(cpu, p.first);
-
-    // Buddy Allocator doesn't manage too large memory chunk
-    if (paddle::memory::Used(cpu) == total_size) continue;
-
-    size_t aligned_size = align(p.second, cpu);
-    total_size -= aligned_size;
-    EXPECT_EQ(total_size, paddle::memory::Used(cpu));
-  }
-}
-#endif

paddle/fluid/memory/memcpy.cc

@@ -27,6 +27,8 @@ void Copy<platform::CPUPlace, platform::CPUPlace>(platform::CPUPlace, void* dst,
 }
 
 #ifdef PADDLE_WITH_CUDA
+static constexpr size_t kMaxGpuAsyncCopyBytes = 64 * 1024;  // 64K
+
 template <>
 void Copy<platform::CPUPlace, platform::CUDAPlace>(
     platform::CPUPlace dst_place, void* dst, platform::CUDAPlace src_place,
@@ -36,6 +38,10 @@ void Copy<platform::CPUPlace, platform::CUDAPlace>(
     platform::GpuMemcpyAsync(dst, src, num, cudaMemcpyDeviceToHost, stream);
   } else {
     platform::GpuMemcpySync(dst, src, num, cudaMemcpyDeviceToHost);
+    // FIXME(zjl): do we really need it?
+    if (num <= kMaxGpuAsyncCopyBytes) {
+      cudaStreamSynchronize(0);
+    }
   }
 }
 
@@ -48,6 +54,10 @@ void Copy<platform::CUDAPlace, platform::CPUPlace>(
     platform::GpuMemcpyAsync(dst, src, num, cudaMemcpyHostToDevice, stream);
   } else {
     platform::GpuMemcpySync(dst, src, num, cudaMemcpyHostToDevice);
+    // FIXME(zjl): do we really need it?
+    if (num <= kMaxGpuAsyncCopyBytes) {
+      cudaStreamSynchronize(0);
+    }
   }
 }
 

paddle/fluid/operators/CMakeLists.txt

@@ -72,7 +72,7 @@ set(OPERATOR_DEPS ${OPERATOR_DEPS} ${COMMON_OP_DEPS})
 set(GLOB_OPERATOR_DEPS ${OPERATOR_DEPS} CACHE INTERNAL "Global Op dependencies")
 
 cc_test(gather_test SRCS gather_test.cc DEPS tensor)
-cc_test(scatter_test SRCS scatter_test.cc DEPS tensor)
+cc_test(scatter_test SRCS scatter_test.cc DEPS tensor math_function)
 cc_test(beam_search_decode_op_test SRCS beam_search_decode_op_test.cc DEPS lod_tensor)
 cc_test(beam_search_op_test SRCS beam_search_op_test.cc DEPS lod_tensor beam_search_op)
 cc_test(strided_memcpy_test SRCS strided_memcpy_test.cc DEPS tensor memory)

paddle/fluid/operators/beam_search_op_test.cc

@@ -54,7 +54,8 @@ void CreateInput(LoDTensor* ids, LoDTensor* scores) {
   }
 }
 
-TEST(beam_search_op, run) {
+// It seems that beam_search_op has bugs.
+TEST(DISABLED_beam_search_op, run) {
   CPUPlace place;
   LoDTensor ids, scores;
   CreateInput(&ids, &scores);

paddle/fluid/operators/conv_mkldnn_op.cc

@@ -12,11 +12,11 @@
    See the License for the specific language governing permissions and
    limitations under the License. */
 
+#include "paddle/fluid/framework/data_layout_transform.h"
+#include "paddle/fluid/memory/malloc.h"
 #include "paddle/fluid/operators/conv_op.h"
 #include "paddle/fluid/platform/mkldnn_helper.h"
 
-#include "paddle/fluid/framework/data_layout_transform.h"
-
 namespace paddle {
 namespace operators {
 
@@ -428,8 +428,9 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
                         "same dimension sizes");
 
       if (residual_param->format() != handler.GetDstFormat()) {
-        auto output_data =
-            output->mutable_data<T>(ctx.GetPlace(), handler.GetDstMemorySize());
+        auto output_data = output->mutable_data<T>(
+            ctx.GetPlace(), ::paddle::memory::Allocator::kDefault,
+            handler.GetDstMemorySize());
         auto residual_data_tz =
             paddle::framework::vectorize2int(residual_param->dims());
         auto residual_data_type =
@@ -449,8 +450,9 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
             handler.AcquireDstMemoryFromPrimitive(to_void_cast<T>(output_data));
       }
     } else {
-      auto output_data =
-          output->mutable_data<T>(ctx.GetPlace(), handler.GetDstMemorySize());
+      auto output_data = output->mutable_data<T>(
+          ctx.GetPlace(), paddle::memory::Allocator::kDefault,
+          handler.GetDstMemorySize());
       dst_memory_p =
           handler.AcquireDstMemoryFromPrimitive(to_void_cast<T>(output_data));
     }
@@ -692,7 +694,8 @@ class ConvMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
               user_diff_dst_memory_p, pipeline);
 
       const size_t size = handler.GetDiffWeightsMemorySize();
-      filter_grad_data = filter_grad->mutable_data<T>(ctx.GetPlace(), size);
+      filter_grad_data = filter_grad->mutable_data<T>(
+          ctx.GetPlace(), paddle::memory::Allocator::kDefault, size);
 
       auto diff_weights_memory_p =
           handler.AcquireDiffWeightsMemoryFromWeightsPrimitive(
@@ -717,7 +720,8 @@ class ConvMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
                                                         pipeline);
 
       const size_t size = handler.GetDiffSourceMemorySize();
-      input_grad_data = input_grad->mutable_data<T>(ctx.GetPlace(), size);
+      input_grad_data = input_grad->mutable_data<T>(
+          ctx.GetPlace(), paddle::memory::Allocator::kDefault, size);
 
       auto diff_src_memory_p = handler.AcquireDiffSrcMemoryFromDataPrimitive(
           reinterpret_cast<void*>(input_grad_data));

paddle/fluid/operators/distributed/variable_response.cc

@@ -115,11 +115,11 @@ bool VariableResponse::CopyLodTensorData(
 
   void* tensor_data =
       tensor->mutable_data(ctx.GetPlace(), ToTypeIndex(meta_.data_type()));
-  if (!ReadRaw(input, ctx, tensor->place(), tensor_data, length)) {
-    return false;
-  }
 
-  return true;
+  VLOG(6) << "Tensor.memory_size = " << tensor->memory_size()
+          << ", Buffer Size = " << length;
+  PADDLE_ENFORCE_EQ(tensor->memory_size(), length);
+  return ReadRaw(input, ctx, tensor->place(), tensor_data, length);
 }
 
 inline framework::DDim GetDims(

paddle/fluid/platform/CMakeLists.txt

@@ -73,3 +73,4 @@ cc_test(float16_test SRCS float16_test.cc DEPS lod_tensor)
 IF(WITH_GPU)
   nv_test(cuda_helper_test SRCS cuda_helper_test.cu)
 ENDIF()
+nv_library(cuda_device_guard SRCS cuda_device_guard.cc DEPS gpu_info)