Realloc for forward

paddle/fluid/operators/activation_mkldnn_op.cc

@@ -23,6 +23,13 @@ using paddle::framework::Tensor;
 using paddle::platform::MKLDNNDeviceContext;
 
 namespace {
+std::string gethash(const mkldnn::memory::dims &operand_dims,
+                    const mkldnn::algorithm algorithm) {
+  return std::string(std::to_string(operand_dims[0]) + "-" +
+                     std::to_string(operand_dims[1]) + "-" +
+                     std::to_string(algorithm));
+}
+
 template <typename T, typename ExecContext>
 void eltwise_forward(const ExecContext &ctx, mkldnn::algorithm algorithm,
                      const T alpha = 0, const T beta = 0) {
@@ -44,6 +51,16 @@ void eltwise_forward(const ExecContext &ctx, mkldnn::algorithm algorithm,
                  "Input dim must be with 2 or 4");
   std::vector<int> src_tz = framework::vectorize2int(src->dims());
 
+  const std::string key = gethash(src_tz, algorithm);
+  const std::string key_src_mem = key + "@eltwise_src_mem";
+  const std::string key_dst_mem = key + "@eltwise_dst_mem";
+  const std::string key_fwd = key + "@eltwise_fwd";
+
+  std::shared_ptr<void> p_src_mem = dev_ctx.GetBlob(key_src_mem);
+  std::shared_ptr<void> p_dst_mem = dev_ctx.GetBlob(key_dst_mem);
+  std::shared_ptr<void> p_fwd = dev_ctx.GetBlob(key_fwd);
+
+  if (p_src_mem == nullptr || p_dst_mem == nullptr || p_fwd == nullptr) {
     // create memory description
     auto data_md = src_tz.size() == 2
                        ? platform::MKLDNNMemDesc(src_tz, mkldnn::memory::f32,
@@ -52,29 +69,35 @@ void eltwise_forward(const ExecContext &ctx, mkldnn::algorithm algorithm,
                                                  mkldnn::memory::format::nchw);
 
     // create memory primitives
-  auto src_memory = std::make_shared<mkldnn::memory>(
+    p_src_mem = std::make_shared<mkldnn::memory>(
         mkldnn::memory({data_md, mkldnn_engine},
                        static_cast<void *>(const_cast<float *>(src_data))));
-  // save source memory to device context to be referred in backward path
-  dev_ctx.SetBlob("InputX@eltwise_pd", src_memory);
-  auto dst_memory =
+    dev_ctx.SetBlob(key_src_mem, p_src_mem);
+
+    p_dst_mem = std::make_shared<mkldnn::memory>(
         mkldnn::memory({data_md, mkldnn_engine},
-                     static_cast<void *>(const_cast<float *>(dst_data)));
+                       static_cast<void *>(const_cast<float *>(dst_data))));
+    dev_ctx.SetBlob(key_dst_mem, p_dst_mem);
 
-  auto forward_desc = mkldnn::eltwise_forward::desc(
+    auto fwd_desc = mkldnn::eltwise_forward::desc(
         mkldnn::prop_kind::forward_training, algorithm, data_md, alpha, beta);
-
-  // save prim desc into global device context to be referred in backward path
-  const std::string key = ctx.op().Output("Out");
-  const std::string key_eltwise_pd = key + "@eltwise_pd";
-  auto forward_pd = std::make_shared<mkldnn::eltwise_forward::primitive_desc>(
-      forward_desc, mkldnn_engine);
-  dev_ctx.SetBlob(key_eltwise_pd, forward_pd);
-
-  auto eltwise = mkldnn::eltwise_forward(*forward_pd, *src_memory, dst_memory);
+    auto p_fwd_pd = std::make_shared<mkldnn::eltwise_forward::primitive_desc>(
+        fwd_desc, mkldnn_engine);
+    p_fwd = std::make_shared<mkldnn::eltwise_forward>(
+        *(p_fwd_pd.get()), *(static_cast<mkldnn::memory *>(p_src_mem.get())),
+        *(static_cast<mkldnn::memory *>(p_dst_mem.get())));
+    dev_ctx.SetBlob(key_fwd, p_fwd);
+  } else {
+    std::static_pointer_cast<mkldnn::memory>(p_src_mem)->set_data_handle(
+        reinterpret_cast<void *>(const_cast<T *>(src_data)));
+
+    std::static_pointer_cast<mkldnn::memory>(p_dst_mem)->set_data_handle(
+        reinterpret_cast<void *>(const_cast<T *>(dst_data)));
+  }
 
   // push primitive to stream and wait until it's executed
-  std::vector<mkldnn::primitive> pipeline = {eltwise};
+  std::vector<mkldnn::primitive> pipeline = {
+      *(static_cast<mkldnn::eltwise_forward::primitive *>(p_fwd.get()))};
   mkldnn::stream(mkldnn::stream::kind::eager).submit(pipeline).wait();
 }
 
@@ -85,7 +108,7 @@ void eltwise_grad(const ExecContext &ctx, mkldnn::algorithm algorithm,
   const auto &mkldnn_engine = dev_ctx.GetEngine();
 
   // get buffers
-  const auto *x = ctx.template Input<Tensor>("Out");
+  const auto *out = ctx.template Input<Tensor>("Out");
 
   auto *dout = ctx.template Input<Tensor>(framework::GradVarName("Out"));
   const auto *diff_dst = dout->template data<T>();
@@ -95,7 +118,12 @@ void eltwise_grad(const ExecContext &ctx, mkldnn::algorithm algorithm,
   const T *diff_src = dx->template mutable_data<T>(ctx.GetPlace());
 
   // get memory dim
-  std::vector<int> src_tz = framework::vectorize2int(x->dims());
+  std::vector<int> src_tz = framework::vectorize2int(out->dims());
+
+  const std::string key = gethash(src_tz, algorithm);
+  const std::string key_src_mem = key + "@eltwise_src_mem";
+  const std::string key_dst_mem = key + "@eltwise_dst_mem";
+  const std::string key_fwd = key + "@eltwise_fwd";
 
   // create memory description
   auto data_md = src_tz.size() == 2
@@ -105,8 +133,8 @@ void eltwise_grad(const ExecContext &ctx, mkldnn::algorithm algorithm,
                                                mkldnn::memory::format::nchw);
 
   // retrieve source memory from device context
-  const std::shared_ptr<void> src_memory = dev_ctx.GetBlob("InputX@eltwise_pd");
-  auto *p_src_memory = static_cast<mkldnn::memory *>(src_memory.get());
+  const std::shared_ptr<void> src_mem = dev_ctx.GetBlob(key_src_mem);
+  auto *p_src_mem = static_cast<mkldnn::memory *>(src_mem.get());
 
   // create memory primitives
   auto diff_src_memory =
@@ -120,9 +148,7 @@ void eltwise_grad(const ExecContext &ctx, mkldnn::algorithm algorithm,
       mkldnn::eltwise_backward::desc(algorithm, data_md, data_md, alpha, beta);
 
   // retrieve eltwise primitive desc from device context
-  const std::string key = ctx.op().Input("Out");
-  const std::string key_eltwise_pd = key + "@eltwise_pd";
-  const std::shared_ptr<void> forward_pd = dev_ctx.GetBlob(key_eltwise_pd);
+  const std::shared_ptr<void> forward_pd = dev_ctx.GetBlob(key_fwd);
   PADDLE_ENFORCE(forward_pd != nullptr,
                  "Fail to find eltwise_pd in device context");
   auto *p_forward_pd =
@@ -131,8 +157,8 @@ void eltwise_grad(const ExecContext &ctx, mkldnn::algorithm algorithm,
   auto eltwise_bwd_prim_desc = mkldnn::eltwise_backward::primitive_desc(
       backward_desc, mkldnn_engine, *p_forward_pd);
 
-  auto eltwise_bwd = mkldnn::eltwise_backward(
-      eltwise_bwd_prim_desc, *p_src_memory, diff_dst_memory, diff_src_memory);
+  auto eltwise_bwd = mkldnn::eltwise_backward(eltwise_bwd_prim_desc, *p_src_mem,
+                                              diff_dst_memory, diff_src_memory);
 
   // push primitive to stream and wait until it's executed
   std::vector<mkldnn::primitive> pipeline = {eltwise_bwd};