From 0cc25a400b1a968f3500273727818af5ee7dc340 Mon Sep 17 00:00:00 2001 From: Krzysztof Binias Date: Wed, 16 May 2018 10:18:41 +0200 Subject: [PATCH] Realloc for forward --- .../fluid/operators/activation_mkldnn_op.cc | 102 +++++++++++------- 1 file changed, 64 insertions(+), 38 deletions(-) diff --git a/paddle/fluid/operators/activation_mkldnn_op.cc b/paddle/fluid/operators/activation_mkldnn_op.cc index fcc06a7093..6f29192550 100644 --- a/paddle/fluid/operators/activation_mkldnn_op.cc +++ b/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 void eltwise_forward(const ExecContext &ctx, mkldnn::algorithm algorithm, const T alpha = 0, const T beta = 0) { @@ -44,37 +51,53 @@ void eltwise_forward(const ExecContext &ctx, mkldnn::algorithm algorithm, "Input dim must be with 2 or 4"); std::vector src_tz = framework::vectorize2int(src->dims()); - // create memory description - auto data_md = src_tz.size() == 2 - ? platform::MKLDNNMemDesc(src_tz, mkldnn::memory::f32, - mkldnn::memory::format::nc) - : platform::MKLDNNMemDesc(src_tz, mkldnn::memory::f32, - mkldnn::memory::format::nchw); - - // create memory primitives - auto src_memory = std::make_shared( - mkldnn::memory({data_md, mkldnn_engine}, - static_cast(const_cast(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 = - mkldnn::memory({data_md, mkldnn_engine}, - static_cast(const_cast(dst_data))); - - auto forward_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( - forward_desc, mkldnn_engine); - dev_ctx.SetBlob(key_eltwise_pd, forward_pd); - - auto eltwise = mkldnn::eltwise_forward(*forward_pd, *src_memory, dst_memory); + 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 p_src_mem = dev_ctx.GetBlob(key_src_mem); + std::shared_ptr p_dst_mem = dev_ctx.GetBlob(key_dst_mem); + std::shared_ptr 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, + mkldnn::memory::format::nc) + : platform::MKLDNNMemDesc(src_tz, mkldnn::memory::f32, + mkldnn::memory::format::nchw); + + // create memory primitives + p_src_mem = std::make_shared( + mkldnn::memory({data_md, mkldnn_engine}, + static_cast(const_cast(src_data)))); + dev_ctx.SetBlob(key_src_mem, p_src_mem); + + p_dst_mem = std::make_shared( + mkldnn::memory({data_md, mkldnn_engine}, + static_cast(const_cast(dst_data)))); + dev_ctx.SetBlob(key_dst_mem, p_dst_mem); + + auto fwd_desc = mkldnn::eltwise_forward::desc( + mkldnn::prop_kind::forward_training, algorithm, data_md, alpha, beta); + auto p_fwd_pd = std::make_shared( + fwd_desc, mkldnn_engine); + p_fwd = std::make_shared( + *(p_fwd_pd.get()), *(static_cast(p_src_mem.get())), + *(static_cast(p_dst_mem.get()))); + dev_ctx.SetBlob(key_fwd, p_fwd); + } else { + std::static_pointer_cast(p_src_mem)->set_data_handle( + reinterpret_cast(const_cast(src_data))); + + std::static_pointer_cast(p_dst_mem)->set_data_handle( + reinterpret_cast(const_cast(dst_data))); + } // push primitive to stream and wait until it's executed - std::vector pipeline = {eltwise}; + std::vector pipeline = { + *(static_cast(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("Out"); + const auto *out = ctx.template Input("Out"); auto *dout = ctx.template Input(framework::GradVarName("Out")); const auto *diff_dst = dout->template data(); @@ -95,7 +118,12 @@ void eltwise_grad(const ExecContext &ctx, mkldnn::algorithm algorithm, const T *diff_src = dx->template mutable_data(ctx.GetPlace()); // get memory dim - std::vector src_tz = framework::vectorize2int(x->dims()); + std::vector 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 src_memory = dev_ctx.GetBlob("InputX@eltwise_pd"); - auto *p_src_memory = static_cast(src_memory.get()); + const std::shared_ptr src_mem = dev_ctx.GetBlob(key_src_mem); + auto *p_src_mem = static_cast(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 forward_pd = dev_ctx.GetBlob(key_eltwise_pd); + const std::shared_ptr 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 pipeline = {eltwise_bwd}; -- GitLab