Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[INTEL MKL] Enabled reorder primitive reuse for conv backprop #20636

Merged
merged 1 commit into from
Jul 20, 2018
Merged

[INTEL MKL] Enabled reorder primitive reuse for conv backprop #20636

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
10 changes: 3 additions & 7 deletions tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -872,12 +872,11 @@ class MklConv2DCustomBackpropFilterOp
}

// check if src and diff_dst need reorder
std::vector<primitive> net;
T *src_data = nullptr;
if (fwd_src_md.data.format != conv2d_bwd_filter->GetSrcMemoryFormat()) {
src.SetUsrMem(fwd_src_md, &src_tensor);
src.CheckReorderToOpMem(
bwd_filter_pd->src_primitive_desc(), &net);
bwd_filter_pd->src_primitive_desc());
src_data = static_cast<T*>(src.GetOpMem().get_data_handle());
} else {
src_data = static_cast<T*>(const_cast<T*>(
Expand All @@ -889,14 +888,13 @@ class MklConv2DCustomBackpropFilterOp
conv2d_bwd_filter->GetDiffDstMemoryFormat()) {
diff_dst.SetUsrMem(diff_dst_md, &diff_dst_tensor);
diff_dst.CheckReorderToOpMem(
bwd_filter_pd->diff_dst_primitive_desc(), &net);
bwd_filter_pd->diff_dst_primitive_desc());
diff_dst_data = static_cast<T*>(
diff_dst.GetOpMem().get_data_handle());
} else {
diff_dst_data = static_cast<T*>(const_cast<T*>(
diff_dst_tensor.flat<T>().data()));
}
stream(stream::kind::eager).submit(net).wait();

// For backward filter, convert diff_filter back to Tensorflow layout
// Here we prepare to reorder op memory back to user memory
Expand Down Expand Up @@ -929,9 +927,7 @@ class MklConv2DCustomBackpropFilterOp

// Reorder diff_filter back to Tensorflow layout if necessary
if (diff_filter_reorder_required) {
std::vector<primitive> net;
diff_filter.InsertReorderToUserMem(&net);
stream(stream::kind::eager).submit(net).wait();
diff_filter.InsertReorderToUserMem();
}
} catch (mkldnn::error& e) {
string error_msg = "Status: " + std::to_string(e.status) +
Expand Down
7 changes: 2 additions & 5 deletions tensorflow/core/kernels/mkl_conv_grad_input_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -722,14 +722,12 @@ class MklConv2DCustomBackpropInputOp
diff_src_tensor->flat<T>().data()));

// check if filter and diff_dst need reorder
std::vector<primitive> net;
T* filter_data = nullptr;
if (fwd_filter_md.data.format !=
conv2d_bwd_input->GetFilterMemoryFormat()) {
filter.SetUsrMem(fwd_filter_md, &filter_tensor);
filter.CheckReorderToOpMem(
bwd_input_pd->weights_primitive_desc(),
&net);
bwd_input_pd->weights_primitive_desc());
filter_data = static_cast<T*>(filter.GetOpMem().get_data_handle());
} else {
filter_data = static_cast<T*>(const_cast<T*>(
Expand All @@ -741,14 +739,13 @@ class MklConv2DCustomBackpropInputOp
conv2d_bwd_input->GetDiffDstMemoryFormat()) {
diff_dst.SetUsrMem(diff_dst_md, &diff_dst_tensor);
diff_dst.CheckReorderToOpMem(
bwd_input_pd->diff_dst_primitive_desc(), &net);
bwd_input_pd->diff_dst_primitive_desc());
diff_dst_data = static_cast<T*>(
diff_dst.GetOpMem().get_data_handle());
} else {
diff_dst_data = static_cast<T*>(const_cast<T*>(
diff_dst_tensor.flat<T>().data()));
}
stream(stream::kind::eager).submit(net).wait();

// execute convolution input bwd
conv2d_bwd_input->Execute(diff_src_data, filter_data, diff_dst_data);
Expand Down