[GitHub] TaoLv commented on a change in pull request #11262: [MXNET-542] Fix mkldnn performance regression + improve test logging

[GitBox <gi...@apache.org>]

TaoLv commented on a change in pull request #11262: [MXNET-542] Fix mkldnn performance regression + improve test logging
URL: https://github.com/apache/incubator-mxnet/pull/11262#discussion_r195891865
 
 

 ##########
 File path: src/operator/nn/mkldnn/mkldnn_base.cc
 ##########
 @@ -77,26 +77,35 @@ mkldnn::memory *TmpMemMgr::Alloc(const mkldnn::memory::primitive_desc &pd) {
   }
 }
 
-mkldnn_output_t CreateMKLDNNMem(const NDArray &arr,
-                                const mkldnn::memory::primitive_desc &desc,
-                                OpReqType req) {
+bool CanWriteTo(const NDArray &out_arr,
+                const NDArray* in_arr,
+                const mkldnn::memory::primitive_desc &desc) {
+  bool add_same = in_arr->GetMKLDNNData()->get_data_handle() ==
+      out_arr.GetMKLDNNData()->get_data_handle();
+  bool pdesc_same = out_arr.GetMKLDNNData()->get_primitive_desc() == desc;
+  return add_same && pdesc_same;
+}
+
+mkldnn_output_t CreateMKLDNNMem(const NDArray &out_arr,
+                                 const mkldnn::memory::primitive_desc &desc,
+                                 OpReqType req, const NDArray* in_arr) {
   if (kAddTo == req) {
     auto tmp = TmpMemMgr::Get()->Alloc(desc);
     return mkldnn_output_t(OutDataOp::AddBack, tmp);
-  } else if (kWriteInplace == req) {
-    // MKLDNN ops may not support the case that the input and the output uses
-    // the same memory. Let's use an extra copy to make sure it always works.
-    auto tmp = TmpMemMgr::Get()->Alloc(desc);
-    return mkldnn_output_t(OutDataOp::CopyBack, tmp);
-  } else {
-    mkldnn::memory *mem = const_cast<NDArray &>(arr).CreateMKLDNNData(desc);
-    if (mem == nullptr) {
-      auto tmp = TmpMemMgr::Get()->Alloc(desc);
-      return mkldnn_output_t(OutDataOp::CopyBack, tmp);
-    } else {
+  } else if (req == kWriteInplace) {
+    if (in_arr != nullptr && CanWriteTo(out_arr, in_arr, desc)) {
+      mkldnn::memory *mem = const_cast<NDArray &>(out_arr).CreateMKLDNNData(desc);
 
 Review comment:
   Why `CreateMKLDNNData` here? I guess in `CanWriteTo` function, you have already called `GetMKLDNNData` from `out_arr`.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services