m53297601
/
Paddle
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Realloc for forward

Browse Source
revert-10836-update-becbos-url
Krzysztof Binias 7 years ago
parent a76d0dd488
commit 0cc25a400b
1 changed files with 64 additions and 38 deletions
Show Stats Download Patch File Download Diff File

76
paddle/fluid/operators/activation_mkldnn_op.cc
Unescape View File

@ -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};

Write Preview
Loading…
Cancel
Save