[GitHub] zheng-da commented on a change in pull request #11308: [MXNET-551] Test CreateMKLDNNMem/CommitOutput

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

zheng-da commented on a change in pull request #11308: [MXNET-551] Test CreateMKLDNNMem/CommitOutput
URL: https://github.com/apache/incubator-mxnet/pull/11308#discussion_r197243674
 
 

 ##########
 File path: src/operator/nn/mkldnn/mkldnn_sum.cc
 ##########
 @@ -36,12 +36,22 @@ void MKLDNNSum(const mkldnn::memory &arr1, const mkldnn::memory &arr2,
   std::vector<mkldnn::memory::primitive_desc> input_pds(2);
   std::vector<float> scales(2, 1);
   std::vector<mkldnn::primitive::at> inputs;
-  input_pds[0] = arr1.get_primitive_desc();
-  input_pds[1] = arr2.get_primitive_desc();
-  CHECK(input_pds[0] == input_pds[1]);
-  inputs.push_back(arr1);
-  inputs.push_back(arr2);
-  // TODO(zhengda) I need to reorder memory here.
+  CHECK(arr1.get_primitive_desc() == arr2.get_primitive_desc());
+  const mkldnn::memory *in_mem1 = &arr1;
+  const mkldnn::memory *in_mem2 = &arr2;
+  if (arr1.get_primitive_desc() != out.get_primitive_desc()) {
+    auto tmp_memory1 = TmpMemMgr::Get()->Alloc(out.get_primitive_desc());
+    auto tmp_memory2 = TmpMemMgr::Get()->Alloc(out.get_primitive_desc());
+    mxnet::MKLDNNCopy(*in_mem1, tmp_memory1);
+    mxnet::MKLDNNCopy(*in_mem2, tmp_memory2);
+    in_mem1 = tmp_memory1;
+    in_mem2 = tmp_memory2;
+  }
+
+  input_pds[0] = in_mem1->get_primitive_desc();
+  input_pds[1] = in_mem2->get_primitive_desc();
 
 Review comment:
   please move these two lines above line 39.
   i found that get_primitive_desc() is quite expensive. we should reduce the number of calls if possible.

----------------------------------------------------------------
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